From 656550c469f587927e5911da1f599c547a9435a2 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 7 Nov 2026 16:01:28 +0800 Subject: [PATCH 01/54] feat: add ReplaySSM state modes for GDN and KDA --- .../attention_state/bounded_att/__init__.py | 10 +- .../attention_state/bounded_att/base_state.py | 4 + .../bounded_att/conv_ssm_state.py | 561 +++++++++++++ .../bounded_att/gdn/__init__.py | 11 +- .../bounded_att/gdn/create_gdn_state.py | 18 +- .../bounded_att/gdn/gdn_state.py | 523 +------------ .../attention_state/bounded_att/interface.py | 15 +- .../bounded_att/kda/__init__.py | 11 +- .../bounded_att/kda/create_kda_state.py | 18 +- .../bounded_att/kda/kda_state.py | 523 +------------ .../sliding_window/fp8_e8m0_swa_state.py | 2 +- .../bounded_att/ssm/__init__.py | 41 + .../attention_state/bounded_att/ssm/base.py | 107 +++ .../bounded_att/ssm/compact.py | 192 +++++ .../attention_state/bounded_att/ssm/export.py | 95 +++ .../bounded_att/ssm/forward.py | 86 ++ .../attention_state/bounded_att/ssm/native.py | 117 +++ .../attention_state/bounded_att/ssm/replay.py | 345 ++++++++ .../window_compressor_state.py | 2 +- lightllm/common/attention_state/commit.py | 70 ++ lightllm/common/attention_state/manager.py | 19 +- .../common/basemodel/attention/base_att.py | 11 + .../basemodel/attention/linear/flashqla.py | 4 +- .../common/basemodel/attention/linear/gdn.py | 233 ++---- .../common/basemodel/attention/linear/kda.py | 148 +--- .../basemodel/attention/linear/state.py | 94 +++ .../basemodel/attention/linear/triton.py | 4 +- lightllm/common/basemodel/basemodel.py | 7 + lightllm/common/basemodel/batch_objs.py | 8 +- lightllm/common/basemodel/cuda_graph.py | 14 +- .../triton_kernel/linear_att/compact_ssm.py | 112 +++ .../linear_att/mtp_state_params.py | 49 +- .../triton_kernel/linear_att/replayssm.py | 31 + .../linear_att/replayssm_common.py | 22 + .../linear_att/replayssm_compute.py | 618 +++++++++++++++ .../linear_att/replayssm_fold.py | 174 +++++ .../linear_att/replayssm_materialize.py | 73 ++ .../linear_att/ssm_acceptance.py | 46 ++ .../triton_kernel/linear_att/ssm_autotune.py | 187 +++++ lightllm/common/ssm_types.py | 13 + lightllm/common/triton_utils/autotuner.py | 20 +- lightllm/models/glm5_next/config.py | 2 + lightllm/models/glm5_next/model.py | 3 + lightllm/models/qwen3next/model.py | 2 + lightllm/server/api_cli.py | 15 + lightllm/server/api_start.py | 65 ++ lightllm/server/core/objs/start_args_type.py | 2 + .../mode_backend/chunked_prefill/impl.py | 3 +- .../mode_backend/dp_backend/impl.py | 8 +- .../model_infer/mtp_speculative/utils.py | 10 +- test/att_resources/validate_model_cache.py | 12 +- test/benchmark/kernels/benchmark_kda_ssm.py | 155 ++++ .../kernels/benchmark_replayssm_service.py | 119 +++ .../common/att_resources/model_helpers.py | 5 + .../att_resources/test_backend_startup.py | 13 +- .../att_resources/test_bounded_state.py | 3 +- .../test_hybrid_layer_mapping.py | 14 +- .../att_resources/test_mtp_state_commit.py | 58 +- .../basemodel/attention/linear/test_gdn.py | 42 +- .../basemodel/test_cuda_graph_autotune.py | 71 ++ .../linear_att/test_acceptance_fusion.py | 223 ++++++ .../linear_att/test_mtp_state_params.py | 68 ++ .../linear_att/test_replayssm.py | 738 ++++++++++++++++++ .../linear_att/test_replayssm_execution.py | 400 ++++++++++ .../linear_att/test_replayssm_grouped_keys.py | 81 ++ .../linear_att/test_replayssm_parallel.py | 172 ++++ .../linear_att/test_ssm_autotune.py | 423 ++++++++++ .../linear_att/test_ssm_autotune_memory.py | 70 ++ unit_tests/common/test_autotuner.py | 58 ++ .../test_hybrid_req_manager_pd_sharing.py | 132 ++++ unit_tests/common/test_linear_state.py | 231 ++++++ .../common/test_linear_state_contract.py | 338 ++++++++ unit_tests/conftest.py | 13 + unit_tests/models/glm5_next/test_cache.py | 23 +- unit_tests/models/glm5_next/test_mtp.py | 57 +- unit_tests/models/glm5_next/test_pd_cache.py | 36 +- .../models/sensenova_loop/test_gdn_state.py | 150 ++++ unit_tests/models/test_post_layer_output.py | 7 +- unit_tests/server/test_hybrid_start_args.py | 30 + unit_tests/server/test_ssm_state_mode.py | 219 ++++++ 80 files changed, 7265 insertions(+), 1444 deletions(-) create mode 100644 lightllm/common/attention_state/bounded_att/conv_ssm_state.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/__init__.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/base.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/compact.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/export.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/forward.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/native.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/replay.py create mode 100644 lightllm/common/attention_state/commit.py create mode 100644 lightllm/common/basemodel/attention/linear/state.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py create mode 100644 lightllm/common/ssm_types.py create mode 100644 test/benchmark/kernels/benchmark_kda_ssm.py create mode 100644 test/benchmark/kernels/benchmark_replayssm_service.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py create mode 100644 unit_tests/common/test_hybrid_req_manager_pd_sharing.py create mode 100644 unit_tests/common/test_linear_state.py create mode 100644 unit_tests/common/test_linear_state_contract.py create mode 100644 unit_tests/models/sensenova_loop/test_gdn_state.py create mode 100644 unit_tests/server/test_ssm_state_mode.py diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py index 55190bea6d..9e5633805e 100644 --- a/lightllm/common/attention_state/bounded_att/__init__.py +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -2,9 +2,9 @@ from .kda.create_kda_state import create_kda_state from .conv_ssm_types import ConvSSMStateConfig from .gdn.gdn_pages_pool import GDNPagesPool -from .gdn.gdn_state import GDNState +from .gdn.gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState from .kda.kda_pages_pool import KDAPagesPool -from .kda.kda_state import KDAState +from .kda.kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState from .base_state import BoundedAttentionState from .sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState @@ -16,9 +16,15 @@ "BoundedAttentionState", "ConvSSMStateConfig", "GDNState", + "GDNNativeState", + "GDNCompactState", + "GDNReplayState", "GDNPagesPool", "create_gdn_state", "KDAState", + "KDANativeState", + "KDACompactState", + "KDAReplayState", "KDAPagesPool", "create_kda_state", "FP8E8M0SlidingWindowConfig", diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 6a96e76dc9..dca8f4e0f3 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, acceptance): + + 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..614e2629de --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -0,0 +1,561 @@ +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_executor +from .ssm.forward import SSMForward +from .ssm.export import export_indexed_state + +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 + state_mode: str | None = None + + 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 + if self.state_mode is not None and ssm_state_mode != self.state_mode: + raise ValueError(f"state mode {ssm_state_mode!r} does not match {self.state_mode!r}") + ssm_state_mode = self.state_mode or ssm_state_mode + 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_executor( + 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, self._ssm, reqs, self._ssm.decode_indices(reqs, speculative), cu_seqlens, accepted) + + def get_accepted_state(self, req_idx): + """Return canonical tensors; native MTP gathers and replay materialization allocate snapshots.""" + 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 export_accepted_state(self, req_indices, targets): + """Write disjoint targets on the current stream without full-state scratch or changing source history.""" + export_indexed_state(self.conv_state, req_indices, [t.conv for t in targets], self.mtp_state_index) + self._ssm.export_accepted_state(req_indices, [t.ssm for t in targets], self.mtp_state_index) + + 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) + self.export_accepted_state( + req_indices, [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；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, acceptance): + assert ( + acceptance.req_indices.device == self.ssm_state.device + ), "commit inputs must match the state device; CUDA states require CUDA tensors" + if self.mtp_state_index is not None: + validator = getattr(self._ssm, "validate_commit", None) + return validator(acceptance, self.state_name) if validator is not None else None + return None + + def commit_bounded_att_state(self, acceptance): + commit_data = self.validate_bounded_att_commit(acceptance) + if self.mtp_state_index is not None: + self._ssm.commit_validated(acceptance, self.mtp_state_index, self.state_name, commit_data=commit_data) + + # ======================================================================== + # 三、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" + remote_dp_bounded_state.export_accepted_state( + src_req_indices, [self.prepare_restore(req) for req in dst_req_indices] + ) + + # ======================================================================== + # 五、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/__init__.py b/lightllm/common/attention_state/bounded_att/gdn/__init__.py index 62dc94c11d..03e545bfa2 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/__init__.py +++ b/lightllm/common/attention_state/bounded_att/gdn/__init__.py @@ -1,5 +1,12 @@ from .create_gdn_state import create_gdn_state from .gdn_pages_pool import GDNPagesPool -from .gdn_state import GDNState +from .gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState -__all__ = ["create_gdn_state", "GDNPagesPool", "GDNState"] +__all__ = [ + "create_gdn_state", + "GDNPagesPool", + "GDNState", + "GDNNativeState", + "GDNCompactState", + "GDNReplayState", +] 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..cfa58d42e9 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .gdn_state import GDNState +from .gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState def create_gdn_state( @@ -13,11 +13,25 @@ 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( + state_cls = { + "native": GDNNativeState, + "compact": GDNCompactState, + "replay": GDNReplayState, + }.get(ssm_state_mode) + if state_cls is None: + raise ValueError(f"unknown GDN state mode: {ssm_state_mode}") + return state_cls( config=config, 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..04c7a52bce 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,29 @@ -from __future__ import annotations +"""GDN 运行态及其三种 SSM 执行策略。""" -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 分离传输 - # ======================================================================== +class GDNState(_ConvSSMState): + """GDN 的公共资源外壳；具体 replay 策略由下列子类固定。""" - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() + algorithm = "gdn" + pages_pool_class = GDNPagesPool - 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) +class GDNNativeState(GDNState): + """GDN 原生 MTP 槽位策略。""" - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== + state_mode = "native" - 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 +class GDNCompactState(GDNState): + """GDN compact 接受前缀重建策略。""" - # 阶段三：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) + state_mode = "compact" - 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 GDNReplayState(GDNState): + """GDN replay history 策略。""" - @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")), - } + state_mode = "replay" diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index fe1fceb3bd..97e6261787 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -8,6 +8,8 @@ import torch +from ..commit import StateAcceptance + from ..utils import deepcopy_result from lightllm.common.attention_state.types import PageKind, StateTransferLayout, TransferRegion @@ -98,12 +100,7 @@ def restore_bounded_att_runtime_from_pages( ... @abstractmethod - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: ... @@ -173,10 +170,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 +199,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/__init__.py b/lightllm/common/attention_state/bounded_att/kda/__init__.py index 7154892202..3fcab8e5d1 100644 --- a/lightllm/common/attention_state/bounded_att/kda/__init__.py +++ b/lightllm/common/attention_state/bounded_att/kda/__init__.py @@ -1,5 +1,12 @@ from .create_kda_state import create_kda_state from .kda_pages_pool import KDAPagesPool -from .kda_state import KDAState +from .kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState -__all__ = ["create_kda_state", "KDAPagesPool", "KDAState"] +__all__ = [ + "create_kda_state", + "KDAPagesPool", + "KDAState", + "KDANativeState", + "KDACompactState", + "KDAReplayState", +] 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..99f7246ea2 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .kda_state import KDAState +from .kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState def create_kda_state( @@ -13,11 +13,25 @@ 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( + state_cls = { + "native": KDANativeState, + "compact": KDACompactState, + "replay": KDAReplayState, + }.get(ssm_state_mode) + if state_cls is None: + raise ValueError(f"unknown KDA state mode: {ssm_state_mode}") + return state_cls( config=config, 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..d7562455f2 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,29 @@ -from __future__ import annotations +"""KDA 运行态及其三种 SSM 执行策略。""" -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 分离传输 - # ======================================================================== +class KDAState(_ConvSSMState): + """KDA 的公共资源外壳；具体 replay 策略由下列子类固定。""" - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() + algorithm = "kda" + pages_pool_class = KDAPagesPool - 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) +class KDANativeState(KDAState): + """KDA 原生 MTP 槽位策略。""" - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== + state_mode = "native" - 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 +class KDACompactState(KDAState): + """KDA compact 接受前缀重建策略。""" - # 阶段三：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) + state_mode = "compact" - 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 KDAReplayState(KDAState): + """KDA replay history 策略。""" - @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")), - } + state_mode = "replay" 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..f6aadf7b61 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,7 @@ 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, acceptance) -> 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..ab61fc5f7d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/__init__.py @@ -0,0 +1,41 @@ +import torch + + +def create_ssm_executor(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound): + """按统一 state layout 创建 native、compact 或 replay executor。 + + 三种模式共享 canonical SSM state；差异只在 decode 中间量和接受阶段的 + state 更新方式。native 为每个 MTP 位置分配物理槽位，compact/replay + 则为每个请求分配一个槽位并在各自 executor 内保存临时数据。 + """ + + 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 NativeSSMExecutor + + return NativeSSMExecutor(state, mtp_step + 1, **kwargs) + kwargs.update(num_key_heads=config.num_k_heads) + if mode == "compact": + from .compact import CompactSSMExecutor + + return CompactSSMExecutor(state, mtp_step + 1, config.conv_dtype, **kwargs) + from .replay import ReplaySSMExecutor + + replay_config = replay_config or {} + return ReplaySSMExecutor( + 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..984a776d26 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -0,0 +1,107 @@ +from lightllm.common.ssm_types import SSMRunConfig + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from .export import export_indexed_state + + +def update_accepted_indices(acceptance, indices, *, replay_cursors=None): + commit_ssm_prefix(indices, acceptance.req_indices, acceptance.accepted_lengths, replay_cursors) + + +class SSMExecutor: + """一种 SSM 模式的 decode、接受和快照适配器。 + + the algorithm-specific state 持有资源和逻辑层到物理层的映射；executor 只处理 + decode 递推、接受前缀以及 accepted-state 导出。prefill 仍由 attention + backend 执行。一次 verify 的布局由 ``SSMForward`` 保存，并在 commit + 阶段通过 ``StateCommit`` 传回同一个 executor。 + """ + + # State 负责资源生命周期和层映射，Executor 只负责本模式的 decode、接受和快照导出。 + # 这样 Native、Compact、Replay 共享同一条上层调用链，差异收敛在执行器内部。 + + slots_per_request = 1 + _run_config: SSMRunConfig | None = None + + @property + def run_config(self) -> SSMRunConfig | None: + """返回最近一次 forward 使用的不可变 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None): + self._run_config = config if config is None or isinstance(config, SSMRunConfig) else SSMRunConfig(**config) + + def clear_history(self, req): + """清空请求的临时 history；默认模式没有额外 history。""" + + return None + + def prepare_decode(self, reqs, cu_seqlens=None): + """在 decode 前准备本轮基准；Native/Compact 默认不需要准备。""" + + return None + + def accept_single_token(self, reqs): + """提交非 speculative decode 产生的单 token 更新。""" + + return None + + def merge_accepted_updates(self, reqs): + """把已接受的临时更新合并回 canonical state。""" + + return None + + def decode_indices(self, reqs, speculative): + """返回 decode 使用的物理槽位；Native MTP 会展开请求槽位。""" + + return reqs + + def get_accepted_state(self, req_idx, accepted_index=None): + """返回一个请求的 accepted-state 视图或物化快照。""" + + return self.state[:, req_idx] + + def export_accepted_state(self, req_indices, targets, accepted_indices=None): + """把 accepted state 写入调用方提供的目标页，不负责分配目标。""" + + export_indexed_state(self.state, req_indices, targets) + + def create_commit(self, run_config): + """把本次 forward 的配置转换为接受阶段可携带的提交 payload。""" + + return None + + def commit(self, acceptance, indices, state_name): + """校验并应用接受结果；供直接调用 State 接口时使用。""" + + validator = getattr(self, "validate_commit", None) + commit_data = validator(acceptance, state_name) if validator is not None else None + self.commit_validated(acceptance, indices, state_name, commit_data=commit_data) + + def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): + """应用已由 State/Manager 校验过的接受结果，不重复检查公共契约。""" + + update_accepted_indices(acceptance, indices) + + 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, + ): + """执行一层 decode，并返回输出及本次实际采用的 run config。""" + + 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..1c1b1d3d46 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/compact.py @@ -0,0 +1,192 @@ +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 SSMExecutor, update_accepted_indices + + +class CompactSSMExecutor(SSMExecutor): + """保存 verify 中间量，并在接受阶段按原 forward 配置重建 state。""" + + 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 = run_config or SSMRunConfig(8, 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: + """选择与 verify layout 一致的 compact kernel 配置。""" + + config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) + return 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, + ): + """执行 compact verify；该模式只接受带 ``cu_seqlens`` 的 speculative 输入。""" + + 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) + _, 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): + """按 accepted index 重建请求前缀，并使用 verify 时的配置保持舍入一致。""" + 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 + 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): + """把本次 verify 的配置作为 compact 接受提交 payload。""" + + return run_config + + def validate_commit(self, acceptance, state_name): + """确认接受结果为每个 microbatch 提供了 compact forward 配置。""" + + if not acceptance.batches: + raise ValueError("compact acceptance requires forward commit batches") + partitions = tuple(acceptance.partitions(state_name)) + for _, config in partitions: + if not isinstance(config, SSMRunConfig): + raise ValueError(f"missing compact forward configuration for {state_name}") + return partitions + + def commit(self, acceptance, indices, state_name): + """先推进通用接受索引，再逐 microbatch 重建 compact state。""" + + commit_data = self.validate_commit(acceptance, state_name) + self.commit_validated(acceptance, indices, state_name, commit_data=commit_data) + + def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): + """使用已验证的 forward 配置完成 compact state 重建。""" + + if commit_data is None: + commit_data = self.validate_commit(acceptance, state_name) + update_accepted_indices(acceptance, indices) + for reqs, config in commit_data: + self.accept_updates(reqs, indices, run_config=config) diff --git a/lightllm/common/attention_state/bounded_att/ssm/export.py b/lightllm/common/attention_state/bounded_att/ssm/export.py new file mode 100644 index 0000000000..f04c4313c5 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/export.py @@ -0,0 +1,95 @@ +from math import prod + +import torch +import triton +import triton.language as tl + +from ..copy_kernel import copy_bounded_att_tensors + + +def upload_indices(values, device, dtype=torch.int64): + """把 Python 索引上传到目标 device，使用 pinned memory 减少同步。""" + + return torch.tensor(values, dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +@triton.jit(do_not_specialize=["count"]) +def _export_indexed( + Source, + Targets, + Reqs, + Accepted, + count, + SHAPE: tl.constexpr, + SOURCE_STRIDES: tl.constexpr, + TARGET_STRIDES: tl.constexpr, + REQUEST_STRIDE: tl.constexpr, + INDEX_STRIDE: tl.constexpr, + SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + # 每个 request 对应一个预分配目标指针；kernel 只做 indexed copy， + # 不改变 source cursor/history，也不负责目标页生命周期。 + total = count.to(tl.int64) * SIZE + for start in range(tl.program_id(0).to(tl.int64) * BLOCK, total, tl.num_programs(0) * BLOCK): + row = start + tl.arange(0, BLOCK) + valid = row < total + item, offset = row // SIZE, row % SIZE + req = tl.load(Reqs + item, valid, 0).to(tl.int64) + accepted = tl.load(Accepted + req, valid, 0).to(tl.int64) + source_offset = req * REQUEST_STRIDE + accepted * INDEX_STRIDE + target_offset = tl.full((BLOCK,), 0, tl.int64) + for axis in tl.static_range(len(SHAPE) - 1, -1, -1): + position = offset % SHAPE[axis] + offset //= SHAPE[axis] + source_offset += position * SOURCE_STRIDES[axis] + target_offset += position * TARGET_STRIDES[axis] + target = tl.load(Targets + item, valid, 0).to(tl.pointer_type(Source.dtype.element_ty)) + values = tl.load(Source + source_offset, valid, 0) + tl.store(target + target_offset, values, valid) + + +def export_indexed_state(source, req_indices, targets, accepted_indices=None, *, slots_per_request=1, index_stride=1): + """将 source 中的 accepted state 写入预分配的 target 页。 + + ``req_indices`` 和 ``targets`` 必须一一对应；MTP 通过 + ``accepted_indices`` 选择请求内槽位。该函数只负责复制，不拥有 target + 的分配和释放生命周期。 + """ + + if not req_indices: + return + assert len(req_indices) == len(targets) + shape, target_strides = targets[0].shape, targets[0].stride() + assert all(t.shape == shape and t.stride() == target_strides and t.dtype == source.dtype for t in targets) + if accepted_indices is None: + copies = [(source[:, req * slots_per_request][..., : shape[-1]], dst) for req, dst in zip(req_indices, targets)] + if source.is_cuda: + copy_bounded_att_tensors(copies) + else: + for src, dst in copies: + dst.copy_(src) + return + if not source.is_cuda: + for req, dst in zip(req_indices, targets): + accepted = int(accepted_indices[req]) + if slots_per_request > 1: + dst.copy_(source[:, req * slots_per_request + accepted]) + else: + dst.copy_(source[:, req, ..., accepted : accepted + shape[-1]]) + return + _export_indexed[(24,)]( + source, + upload_indices([t.data_ptr() for t in targets], source.device), + upload_indices(req_indices, source.device, torch.int32), + accepted_indices, + len(req_indices), + tuple(shape), + (source.stride(0), *source.stride()[2:]), + target_strides, + source.stride(1) * slots_per_request, + index_stride, + prod(shape), + 1024, + num_warps=8, + ) 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..6930630d53 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/forward.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass, field +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 .base import SSMExecutor + from ..conv_ssm_state import _ConvSSMState + + +@dataclass +class SSMForward: + """携带一次 forward 的请求布局，并生成对应的状态提交凭据。 + + ``req_indices`` 是本次 forward 的 request-row 序列，可能包含 CUDA Graph + 的 ``HOLD`` padding；``state_indices`` 是 kernel 使用的物理槽位， + ``cu_seqlens`` 描述真实 query 窗口。``run_config`` 在第一层执行时确定， + 后续层复用同一配置，``finish`` 再把它封装进 ``StateCommit``。 + """ + + resource: "_ConvSSMState" + executor: "SSMExecutor" + req_indices: torch.Tensor + state_indices: torch.Tensor + cu_seqlens: Optional[torch.Tensor] + accepted_token_counts: Optional[torch.Tensor] + run_config: Optional[SSMRunConfig] = None + _layout_checked: bool = field(default=False, init=False, repr=False) + + def run_layer(self, logical_layer, q, k, v, a, b, a_log, bias): + """执行一层线性注意力；logical_layer 会映射到物理权重层。""" + + output, self.run_config = self.executor.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): + """结束 forward，返回接受阶段需要的状态提交描述。""" + + payload = self.executor.create_commit(self.run_config) + if payload is None: + return None + return StateCommit(self.resource.state_name, payload) + + def copy_inputs(self, source: "SSMForward"): + """把动态图输入复制到已 capture 的 tensor。 + + graph bucket 的资源和 tensor layout 在首次绑定时校验一次，后续调用 + 只更新数据；captured layout 不随请求变化。 + """ + + names = ("req_indices", "state_indices", "cu_seqlens", "accepted_token_counts") + if not self._layout_checked: + if self.resource is not source.resource or self.executor is not source.executor: + raise ValueError("cannot copy a forward for a different resource or executor") + for name in names: + target, value = getattr(self, name), getattr(source, name) + if (target is None) != (value is None) or ( + target is not None + and (target.shape != value.shape or target.dtype != value.dtype or target.device != value.device) + ): + raise ValueError(f"captured SSM layout changed: {name}") + self._layout_checked = True + + copied = set() + for name in names: + target, value = getattr(self, name), getattr(source, name) + 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..9303efcc2e --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/native.py @@ -0,0 +1,117 @@ +import torch + +from .base import SSMExecutor +from .export import export_indexed_state + + +class NativeSSMExecutor(SSMExecutor): + """直接调用 fused recurrent kernel，并把 MTP 状态展开为物理槽位。""" + + 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): + """返回普通 decode 或 speculative verify 对应的物理 state 索引。""" + + 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): + """返回单槽位 view，或 MTP 接受位置对应的 state 副本。""" + + 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 export_accepted_state(self, req_indices, targets, accepted_indices=None): + """把 native state 按请求和接受位置导出到调用方目标页。""" + + export_indexed_state( + self.state, + req_indices, + targets, + accepted_indices if self.slots_per_request > 1 else None, + slots_per_request=self.slots_per_request, + index_stride=self.state.stride(1), + ) + + 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, + ): + """执行 native recurrent decode，并原地更新 fused kernel 的最终 state。""" + + 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..629bd35da8 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/replay.py @@ -0,0 +1,345 @@ +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 SSMExecutor, update_accepted_indices +from .export import upload_indices + + +class ReplaySSMExecutor(SSMExecutor): + """以短 history 重放 decode 的 SSM 执行器。 + + ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 + cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode + 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 + checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 + """ + + 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", + ): + """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 + + ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行 + 方式，``fold_method`` 选择 raw 或 derived history 的折叠路径。 + """ + + 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): + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + 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): + """更新 cursor 以保留每个请求本轮已接受的前缀。 + + MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft + 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 + """ + _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): + """把接受的 history 物化回 canonical state，并清空请求 cursor。""" + 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, output_strides=None): + """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + + layers, slots, hv, k, v = self.state.shape + config = self.run_config or SSMRunConfig(32, 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, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def get_accepted_state(self, req_idx, accepted_index=None): + """返回不修改 active state/history 的单请求 canonical 快照。""" + 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 export_accepted_state(self, req_indices, targets, accepted_indices=None): + """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" + + if not req_indices: + return + shape = (self.state.shape[0], *self.state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) + self._materialize_accepted_state( + upload_indices(req_indices, self.state.device, torch.int32), + upload_indices([t.data_ptr() for t in targets], self.state.device), + snapshot=True, + output_strides=strides, + ) + + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" + + if self.compute_mode == "parallel": + config = self.run_config or SSMRunConfig(128, 2, 1) + else: + config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + return 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, + ): + """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" + + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + 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, + 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): + """非 MTP decode 直接把单 token 写入当前 history。""" + self.accept_updates(reqs) + + def commit(self, acceptance, indices, state_name): + """应用接受结果；MTP 直接推进 cursor，单 token 还需写入 history。""" + + self.commit_validated(acceptance, indices, state_name) + + def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): + """应用已验证的接受结果，推进 replay cursor。""" + + if self.verify_width > 1: + update_accepted_indices(acceptance, indices, replay_cursors=self.cursors) + else: + update_accepted_indices(acceptance, indices) + self.accept_updates(acceptance.req_indices, 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..958295b1ba 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,7 @@ 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, acceptance) -> 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..6256ea2c80 --- /dev/null +++ b/lightllm/common/attention_state/commit.py @@ -0,0 +1,70 @@ +from dataclasses import dataclass +from typing import Tuple + +import torch + + +@dataclass(frozen=True) +class StateCommit: + """保存一次 forward 产生的状态提交配置。""" + + state_name: str + payload: object + + +@dataclass(frozen=True) +class StateCommitBatch: + """描述一个 microbatch 分区及其各状态的 forward 配置。""" + + 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 StateAcceptance: + """描述一次 speculative verify 的接受结果。 + + ``req_indices`` 和 ``accepted_lengths`` 按未 padding 的 forward 顺序配对， + 接受长度包含基准输入 token。``batches`` 只保存相同顺序下的 microbatch + 配置，提交时由状态执行器按分区恢复产生本次结果所用的 kernel layout。 + """ + + # 接受结果是状态提交的唯一输入；batches 只携带对应 forward 的执行配置，避免调用方 + # 重新拼装请求顺序、接受长度和 microbatch 分区。 + + req_indices: torch.Tensor + accepted_lengths: torch.Tensor + batches: Tuple[StateCommitBatch, ...] = () + + def validate(self, state_names=None): + """检查请求、接受长度和 forward 分区是否能组成一次提交。""" + + if self.req_indices.ndim != 1 or self.accepted_lengths.shape != self.req_indices.shape: + raise ValueError("request indices and accepted lengths must share one-dimensional layout") + for value in (self.req_indices, self.accepted_lengths): + if value.dtype not in (torch.int32, torch.int64) or value.device != self.req_indices.device: + raise ValueError("acceptance tensors must be integer tensors on the same device") + if self.batches and sum(batch.request_count for batch in self.batches) != self.req_indices.numel(): + 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): + """按 forward 分区返回请求索引和指定状态的提交配置。""" + + 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 self.req_indices[start:end], payload + start = end diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..4125ff727e 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -6,6 +6,8 @@ import torch +from .commit import StateAcceptance + from lightllm.common.attention_state.types import PageKind from lightllm.utils.dist_utils import get_current_device_id @@ -392,17 +394,14 @@ def restore_bounded_att_runtime_from_pages( ) # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - role: str, - ) -> None: - """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + def commit_bounded_att_state(self, acceptance: StateAcceptance, role: str) -> None: + """提交按请求排列的已接受前缀；无参与 State 时为空操作。""" + assert acceptance.req_indices.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) + acceptance.validate({state.state_name for state in states}) + for state in states: + state.commit_bounded_att_state(acceptance) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index ae37385e68..b3aa1c8592 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 @@ -163,6 +164,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..1ce0a534b6 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/state.py @@ -0,0 +1,94 @@ +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) + _layout_checked: bool = field(default=False, init=False, repr=False) + + def get_draft_step(self): + return self.backend.mtp_step + + def init_state(self): + self._layout_checked = False + 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()} + state._layout_checked = False + return state + + def copy_for_decode_cuda_graph(self, new_state): + if not self._layout_checked: + 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]) + self._layout_checked = True 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..4a046ef02d 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -21,6 +21,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 +133,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 +166,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..876cfb9101 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Public ReplaySSM kernel exports. + +Kernel implementations are split by lifecycle stage; these imports preserve the +existing module-level names used by the executor and tests. +""" + +from .replayssm_common import _replay_update +from .replayssm_compute import ( + _project_replay_queries, + _replay, + _replay_gdn_window, + _replay_sequential, + _solve_gdn_deltas, +) +from .replayssm_fold import _accept_updates, _fold_history, _prepare_fold, _use_separate_fold +from .replayssm_materialize import _materialize + +__all__ = [ + "_accept_updates", + "_fold_history", + "_materialize", + "_prepare_fold", + "_project_replay_queries", + "_replay", + "_replay_gdn_window", + "_replay_sequential", + "_replay_update", + "_solve_gdn_deltas", + "_use_separate_fold", +] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py new file mode 100644 index 0000000000..a54b264947 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared ReplaySSM recurrence primitives.""" + +import triton +import triton.language as tl + + +@triton.jit +def _replay_update(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + # 所有 replay fold 路径共用这一条递推，保证 raw/derived/history + # materialize 使用相同的数值顺序。 + 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 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py new file mode 100644 index 0000000000..de216db8a7 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py @@ -0,0 +1,618 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM projection and recurrent decode kernels.""" + +import triton +import triton.language as tl + +from .replayssm_common import _replay_update + + +@triton.jit +def _project_replay_queries( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ: tl.constexpr, + SK: tl.constexpr, + SA: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + L: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + PARALLEL_SOLVE: tl.constexpr, + BQK: tl.constexpr, + base, + end, + hv, + kk, + ll, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + slot, + start, + state, +): + # 预计算 verify window 的 q/k 投影；仅 precompute/parallel 路径调用。 + 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) + return qk, projections, tt, tokens, valid, h + + +@triton.jit +def _solve_gdn_deltas(power, rhs, BW: tl.constexpr, projection_dtype: tl.constexpr, projection_precision: tl.constexpr): + # 解同一 verify window 内的严格下三角 delta 耦合，不读写 history。 + wi = tl.arange(0, power.shape[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") + return tl.dot(inverse.to(projection_dtype), rhs.to(projection_dtype), input_precision=projection_precision) + + +@triton.jit +def _replay_gdn_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + PARALLEL_SOLVE: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, + base, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, +): + # GDN precompute 路径：一次投影整个窗口，再按 token 顺序写入新 history。 + # 只有 corrected-value solve 依赖 token 顺序，避免重复计算历史投影。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + 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) + solved = _solve_gdn_deltas(power, rhs, BW, projection_dtype, 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 + + +@triton.jit +def _replay_sequential( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + base, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, +): + # KDA、inline 以及不满足 precompute 条件的路径使用逐 token 递推。 + 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 _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, + KDA: tl.constexpr, + LOWER: tl.constexpr, + PARALLEL_SOLVE: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, +): + """执行一个请求窗口的 replay decode，并写入 raw/derived history。 + + 当 history 即将跨半区时，先按 raw input 折叠到 canonical state；随后 + 根据 compute mode 选择预计算投影或逐 token 递推。请求数属于运行时布局， + 不应因为 eager batch 改变而形成新的 kernel 变体。 + """ + 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 not PARALLEL_SOLVE 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" + qk, projections, tt, tokens, valid, h = _project_replay_queries( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ, + SK, + SA, + H, + HV, + K, + L, + BW, + KDA, + LOWER, + PARALLEL_SOLVE, + BQK, + base, + end, + hv, + kk, + ll, + n, + projection_dtype, + projection_precision, + slot, + start, + state, + ) + else: + projections, tt = None, None + + if PRECOMPUTE_PROJECTIONS and not KDA: + _replay_gdn_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + PARALLEL_SOLVE, + BQK, + BR, + base, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype, + projection_precision, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, + ) + else: + _replay_sequential( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ, + SK, + SV, + SA, + SB, + H, + HV, + KH, + K, + V, + L, + PRECOMPUTE_PROJECTIONS, + BW, + KDA, + LOWER, + base, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py new file mode 100644 index 0000000000..92a1108b0c --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM history preparation and accepted-prefix folding kernels.""" + +import triton +import triton.language as tl + +from .replayssm_common import _replay_update + + +@triton.jit(do_not_specialize=["N"]) +def _accept_updates( + cursors, + reqs, + accepted_state_indices, + N, + HOLD: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 + 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 _use_separate_fold(valid, length, count, N, L: tl.constexpr, WIDTH: tl.constexpr, FORCE: tl.constexpr): + # 只决定本轮是否单独启动 fold；真正的 history 合并由 _fold_history 完成。 + small_batch: tl.constexpr = 12 + dense_batch: tl.constexpr = 256 + live = tl.sum((valid & (length > 0)).to(tl.int32), axis=0) + split = (count > 0) & ((live <= small_batch) | (2 * count <= live)) + if L == 8 and WIDTH == 4: + if N >= dense_batch: + short = tl.sum((valid & (length > 0) & (length < WIDTH)).to(tl.int32), axis=0) + split = split & ((live < dense_batch) | (short > 0)) + if FORCE: + split = count > 0 + return split + + +@triton.jit(do_not_specialize=["N"]) +def _prepare_fold( + Reqs, + Cursors, + Positions, + ForwardCursors, + Cu, + Active, + N, + HOLD: tl.constexpr, + L: tl.constexpr, + VARLEN: tl.constexpr, + BLOCK: tl.constexpr, + WIDTH: tl.constexpr, + SEPARATE_FOLD: tl.constexpr, +): + # 输出 Positions/ForwardCursors 供 decode 使用，Active/N 供后续 fold + # 使用；只改写本轮请求对应的 cursor，不触碰 canonical state。 + 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) + split = _use_separate_fold(valid, length, count, N, L, WIDTH, SEPARATE_FOLD) + 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(do_not_specialize=["N"]) +def _fold_history( + State, + RawKeys, + RawValues, + Gates, + Betas, + Reqs, + Cursors, + Active, + N, + 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, +): + # 只消费 Active 中的 fold 请求，把 history 合并回 canonical state； + # cursor 的 phase 已由 _prepare_fold 切换，写入后新 history 从同一半区继续。 + 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) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py new file mode 100644 index 0000000000..7f82d9f19e --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM accepted-state materialization kernel.""" + +import triton +import triton.language as tl + + +@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, + OUTPUT_STRIDES: tl.constexpr = None, +): + # 将 cursor 指向的 accepted history 应用到 canonical state；SNAPSHOT + # 只写独立输出，merge 路径则原地更新请求槽位。 + 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, :] + if OUTPUT_STRIDES is not None: + target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) + out = ( + target + + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + + head * OUTPUT_STRIDES[1] + + kk[:, None] * OUTPUT_STRIDES[2] + + vv[None, :] * OUTPUT_STRIDES[3] + ) + else: + 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_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py new file mode 100644 index 0000000000..058a0e6d65 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -0,0 +1,46 @@ +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _commit_prefix( + Indices, + Reqs, + Lengths, + Cursors, + N, + HOLD: tl.constexpr, + REQ_STRIDE: tl.constexpr, + LENGTH_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) + length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) + tl.store(Indices + req, length - 1, row < N) + if Cursors is not None: + valid = (row < N) & (req != HOLD) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor + length, valid) + + +def commit_ssm_prefix(indices, reqs, lengths, replay_cursors=None): + if not reqs.is_cuda: + indices[reqs] = (lengths - 1).to(indices.dtype) + if replay_cursors is not None: + valid = reqs != replay_cursors.numel() - 1 + replay_cursors[reqs[valid]] += lengths[valid].to(replay_cursors.dtype) + return + if reqs.numel(): + _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( + indices, + reqs, + lengths, + replay_cursors, + reqs.numel(), + indices.numel() - 1, + reqs.stride(0), + lengths.stride(0), + 256, + num_warps=4, + ) 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..ac7e920aaa --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -0,0 +1,187 @@ +"""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 +from lightllm.common.ssm_types import SSMRunConfig + +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 SSMRunConfig(32, 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 = SSMRunConfig(**config) if config else 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/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/model.py b/lightllm/models/glm5_next/model.py index e15efe0143..a2a7d10ef7 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/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 41ae769ae6..6911939f9e 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -1,4 +1,5 @@ import argparse +import json def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -984,6 +985,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 a42acef884..302c04d4de 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -43,6 +43,69 @@ 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 + + model_type = get_model_type(args.model_dir) + is_gdn_or_kda = model_type in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "sensenova_loop", + ) + if is_gdn_or_kda: + assert args.mtp_step <= 15, ( + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" + ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + 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" + if config["compute_mode"] == "parallel": + assert model_type not in ("glm5_next", "glm5_next_text"), "Parallel ReplaySSM requires GDN, not KDA" + + def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) _set_envs_and_config(args) @@ -355,6 +418,8 @@ def _launch_subprocesses(args: StartArgs): s3_model_prepare(args.model_dir) + _validate_and_normalize_ssm_config(args) + # 如果args.eos_id 是 None, 从 config.json 中读取 eos_token_id 相关的信息，赋值给 args if args.eos_id is None: from lightllm.utils.config_utils import get_eos_token_ids diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0d6468f397..9cf0f719ab 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -255,6 +255,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + ssm_state_mode: str = "native" + ssm_replay_config: dict = 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/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 9bc8c16cf7..5aae49b199 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,7 @@ import torch import time from typing import List +from lightllm.common.attention_state.commit import StateCommitBatch 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 @@ -295,7 +296,7 @@ def decode_mtp( next_token_ids=next_token_ids, 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..1d54f7a21c 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,7 @@ import torch import time from typing import List, Tuple +from lightllm.common.attention_state.commit import StateCommitBatch 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 @@ -530,7 +531,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): next_token_ids=next_token_ids, 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", @@ -804,7 +805,10 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf next_token_ids=next_token_ids, 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..c32d0a970a 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 StateAcceptance + from lightllm.common.basemodel.triton_kernel.mtp_utils import ( mtp_scatter_next_token_ids, mtp_verify, @@ -21,10 +23,9 @@ def verify_mtp_tokens( next_token_ids: torch.Tensor, 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,9 +33,8 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - backend.model.att_state_manager.commit_bounded_att_state( - b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" - ) + acceptance = StateAcceptance(b_req_idx[b_req_mtp_start_loc.long()], accept_lengths, state_commit_batches) + backend.model.att_state_manager.commit_bounded_att_state(acceptance, role="main") return accept_lengths, accepted_index diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 4bf1514704..12bc610eb4 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -22,6 +22,7 @@ import numpy as np import torch +from lightllm.common.attention_state.commit import StateAcceptance, StateCommitBatch import torch.distributed as dist from transformers import PretrainedConfig @@ -227,19 +228,24 @@ def decode(request: torch.Tensor): max_kv_seq_len=int(lengths[-1]), ) output = model.forward(batch) + commits = [output.state_commits] logits = [output.logits.clone()] for draft in draft_models: batch.input_ids = output.logits.argmax(dim=-1) batch.mtp_draft_input_hiddens = output.mtp_collector.spec_hidden output = draft.forward(batch) + commits.append(output.state_commits) logits.append(output.logits.clone()) if draft_models and hybrid: cuda_req_indices = request.cuda(non_blocking=True) accepted_steps = torch.full_like(cuda_req_indices, accepted_width, dtype=torch.int32) - for item in (model, *draft_models): + for item, state_commits in zip((model, *draft_models), commits): resources.commit_bounded_att_state( - req_indices=cuda_req_indices, - accepted_steps=accepted_steps, + acceptance=StateAcceptance( + cuda_req_indices, + accepted_steps, + (StateCommitBatch(request.numel(), state_commits),), + ), role=item.att_state_role, ) return logits diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..3b702f206d --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,155 @@ +"""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. +""" + +from dataclasses import asdict +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 ReplaySSMExecutor +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor +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 = CompactSSMExecutor(state, width, q.dtype, kda=True) + elif mode != "native": + capacity = max(8, triton.next_power_of_2(width)) + ssm_backend = ReplaySSMExecutor(state, capacity, 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, ReplaySSMExecutor): + 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}]: + if mode == "precompute" and config["BV"] < 16: + continue + 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, + cache_len=getattr(ssm_backend, "capacity", None), + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=asdict(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_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e3a09adb8a..e7265b3f72 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -14,6 +14,7 @@ import pytest import torch +from lightllm.common.attention_state.commit import StateAcceptance from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3next.model import Qwen3NextTpPartModel @@ -296,8 +297,10 @@ def start_runtime(): draft_states = [manager.get_bounded_attention_states(f"draft{i}")[0] for i in range(draft_count)] assert len({state.conv_state.data_ptr() for state in draft_states}) == draft_count manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), + ), role="main", ) assert all(not state.mtp_state_index.any() for state in draft_states) @@ -306,8 +309,10 @@ def start_runtime(): for i, state in enumerate(draft_states): assert len(state.config.logical_layer_ids) == 1 manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([i + 2], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), + ), role=f"draft{i}", ) assert state.mtp_state_index[0] == i + 1 diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index ffd59db4aa..21c08329e5 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -5,6 +5,7 @@ import pytest import torch +from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ( @@ -254,7 +255,7 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): assert (state.small_pages_pool.conv_state[0] == 7).all() assert (state.small_pages_pool.ssm_state[0] == 11).all() with pytest.raises(AssertionError, match="CUDA tensors"): - state.commit_bounded_att_state(torch.tensor([0]), torch.tensor([1], dtype=torch.int32)) + state.commit_bounded_att_state(StateAcceptance(torch.tensor([0]), torch.tensor([1], dtype=torch.int32))) @pytest.mark.parametrize("cls", [KDAState, GDNState]) diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 15d29bd85d..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..1220f3a5f5 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -7,6 +7,7 @@ import pytest import torch +from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import GDNState @@ -34,8 +35,9 @@ def test_commit_selects_state_for_one_role(): state.ssm_state[:, slot].fill_(slot + 10) state.mtp_state_index.fill_(1) manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([3], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ), role="main", ) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.main_gdn_conv_fp32_ssm_fp32_state): @@ -44,8 +46,9 @@ def test_commit_selects_state_for_one_role(): assert (state.ssm_state[:, selected] == 12).all() assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [1, 1, 1] manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([1], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([1], dtype=torch.int32, device="cuda") + ), role="draft0", ) assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 1] @@ -58,19 +61,22 @@ def test_commit_without_recurrent_state_and_invalid_role(): manager = AttentionStateManager(request_capacity=3) manager.finish_state_registration() manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([3], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ), role="main", ) manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([3], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ), role="draft0", ) with pytest.raises(AssertionError, match="unknown model role"): manager.commit_bounded_att_state( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([3], dtype=torch.int32, device="cuda"), + StateAcceptance( + torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ), role="target", ) @@ -80,6 +86,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 +94,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)): @@ -100,10 +107,37 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i next_token_ids=torch.tensor(tokens, device="cuda"), b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), ) assert accept_lengths.tolist() == [3, 1, 2] assert accepted_index.tolist() == expected_mask assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] + + +def test_verify_routes_one_acceptance(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) + 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 acceptance, **kw: calls.append((acceptance, 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, state_commit_batches=batches) + assert len(calls) == 1 + acceptance, kwargs = calls[0] + assert acceptance.req_indices.tolist() == [2, 0] + assert acceptance.accepted_lengths is lengths + assert acceptance.batches is batches + assert kwargs == {"role": "main"} 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..d417381753 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 CompactSSMExecutor + 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..5b0f4b249f --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,223 @@ +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 ReplaySSMExecutor +from types import SimpleNamespace +from lightllm.common.attention_state.commit import StateAcceptance +from lightllm.common.attention_state.bounded_att.ssm.native import NativeSSMExecutor + + +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 = ReplaySSMExecutor.__new__(ReplaySSMExecutor) + 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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("replay", [False, True]) +def test_request_counts_share_compiled_kernel(monkeypatch, replay): + from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + + reqs = torch.arange(512, dtype=torch.int32, device="cuda") + lengths = reqs % 4 + 1 + indices = torch.zeros(513, dtype=torch.int32, device="cuda") + cursors = torch.zeros_like(indices) if replay else None + hashes = set() + run = _commit_prefix.run + + def record(*args, **kwargs): + kernel = run(*args, **kwargs) + hashes.add(kernel.hash) + return kernel + + monkeypatch.setattr(_commit_prefix, "run", record) + for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): + indices.fill_(-1) + if replay: + cursors.zero_() + commit_ssm_prefix(indices, reqs[:count], lengths[:count], cursors) + assert torch.equal(indices[:count], lengths[:count] - 1) + assert (indices[count:] == -1).all() + if replay: + assert torch.equal(cursors[:count], lengths[:count]) + assert not cursors[count:].any() + assert len(hashes) == 1 + + +def accept_unfused(m, args): + update(m.mtp_state_index, *args) + m._ssm.accept_updates(args[1][args[0].long()], m.mtp_state_index) + + +def accept_fused(m, args): + reqs = args[1][args[0].long()] + ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) + cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) + lengths = cumulative[ends.long()] - cumulative[args[0].long()] + m._ssm.commit(StateAcceptance(reqs, lengths), m.mtp_state_index, "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 = NativeSSMExecutor(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 = [ + ReplaySSMExecutor( + 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..ee83837a04 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,738 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor +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 = ReplaySSMExecutor( + 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 = ReplaySSMExecutor( + 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 = [ + ReplaySSMExecutor( + 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 = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + for ssm_backend in ssm_backends: + ssm_backend.raw_keys.zero_() + ssm_backend.raw_values.zero_() + ssm_backend.gates.fill_(-0.01) + ssm_backend.betas.zero_() + ssm_backend.cursors[0] = 16 + ssm_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 = ReplaySSMExecutor( + 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 = ReplaySSMExecutor(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 CompactSSMExecutor + + 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 = CompactSSMExecutor(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 CompactSSMExecutor + 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 = CompactSSMExecutor(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, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + "bucket,live,folding,width,separate", + [ + (12, 12, 12, 4, True), + (13, 13, 13, 4, False), + (255, 255, 127, 4, True), + (256, 256, 128, 4, False), + (257, 257, 128, 4, False), + (256, 256, 128, 3, True), + (256, 12, 12, 4, True), + (256, 13, 4, 4, True), + ], +) +def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): + from lightllm.common.attention_state.commit import StateAcceptance + + torch.manual_seed(72) + initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 + backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) + reqs = torch.cat( + ( + torch.arange(live, device="cuda", dtype=torch.int32), + torch.full((bucket - live,), live, device="cuda", dtype=torch.int32), + ) + ) + indices = torch.zeros(live + 1, device="cuda", dtype=torch.int32) + expected = initial[0, :live].clone() + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(32 if kda else 1, device="cuda") + for step, query_width in enumerate((4, 4, width)): + if step == 2: + backend.clear_history(reqs[folding:live]) + expected[folding:] = initial[0, folding:live] + cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) + q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") + b = torch.randn(live, query_width, 1, device="cuda") + backend.prepare_decode(reqs, cu) + if step == 2: + inline = backend.forward_cursors[:live] % 16 + query_width > 8 + assert inline[:folding].eq(not separate).all() + assert not inline[folding:].any() + actual = backend.forward( + 0, + q.reshape(1, -1, 1, 32), + k.reshape(1, -1, 1, 32), + v.reshape(1, -1, 1, 32), + a.flatten(0, 1), + b.flatten(0, 1), + alog, + bias, + reqs, + cu, + )[0].reshape(live, query_width, 1, 32) + for token in range(query_width): + output, expected = reference( + q[:, token], + k[:, token], + v[:, token], + a[:, token], + b[:, token], + alog, + bias.expand(live, 1, 32) if kda else bias, + expected, + kda=kda, + ) + torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) + lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) + backend.commit(StateAcceptance(reqs[:live], lengths), indices, "main") + backend.merge_accepted_updates(reqs[:live]) + torch.testing.assert_close(backend.state[0, :live], expected, atol=3e-4, rtol=3e-4) + assert torch.equal(backend.state[0, live], initial[0, live]) 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..c3efd52b1f --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,400 @@ +import pytest +import torch +import triton + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor +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 StateAcceptance + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): + from lightllm.common.basemodel.triton_kernel.linear_att import replayssm + + kernels = ("_prepare_fold", "_fold_history", "_replay", "_accept_updates") + hashes = {name: set() for name in kernels} + launches = [] + for name in kernels: + kernel = getattr(replayssm, name) + + def record(*args, run=kernel.run, name=name, **kwargs): + compiled = run(*args, **kwargs) + hashes[name].add(compiled.hash) + launches.append(name) + return compiled + + monkeypatch.setattr(kernel, "run", record) + + torch.manual_seed(137) + initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 + eager, captured = [ + ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) + for _ in range(2) + ] + reqs = torch.arange(512, device="cuda", dtype=torch.int32) + q = torch.randn(512, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((512, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(512, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + + def cycle(backend, ids): + count = ids.numel() + backend.prepare_decode(ids) + out = backend.forward(0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] + backend.accept_single_token(ids) + return out + + graphs = {} + for bucket in (32, 512): + ids = reqs[:bucket].clone() + cycle(captured, ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = cycle(captured, ids) + graphs[bucket] = graph, ids, out + captured.state.copy_(initial) + captured.cursors.zero_() + counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) + for count in counts: + graph, ids, output = graphs[32 if count <= 32 else 512] + ids.fill_(captured.hold_req_id) + ids[:count].copy_(reqs[:count]) + for _ in range(9): + expected = cycle(eager, reqs[:count]) + calls = len(launches) + graph.replay() + assert len(launches) == calls + torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) + assert torch.equal(captured.cursors, eager.cursors) + torch.testing.assert_close(captured.state, eager.state, atol=3e-5, rtol=3e-5) + assert len(hashes["_prepare_fold"]) == len({triton.next_power_of_2(n) for n in counts}) + for name in kernels[1:]: + assert len(hashes[name]) == 1, (name, hashes[name]) + assert captured.cursors[-1] == 0 + assert torch.equal(captured.state[:, -1], initial[:, -1]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): + torch.manual_seed(827) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + cls( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + kda=kda, + compute_mode="precompute", + run_config=config, + ) + for cls in [ReplaySSMExecutor, ReplaySSMExecutor] + ] + 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( + StateAcceptance(rows[starts.long()], accepted[rows[starts.long()]] + 1), + manager.mtp_state_index, + "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 = [ + ReplaySSMExecutor( + 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 = [ + ReplaySSMExecutor(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 = [ + ReplaySSMExecutor(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..04954d4361 --- /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 ReplaySSMExecutor +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor + +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 ReplaySSMExecutor(state.clone(), 16, width, num_key_heads=key_heads) + return CompactSSMExecutor(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..a423697c51 --- /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 ReplaySSMExecutor + +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 = ReplaySSMExecutor(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"): + ReplaySSMExecutor(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): + ReplaySSMExecutor(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 = ReplaySSMExecutor( + 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..4ffbcfc937 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -0,0 +1,423 @@ +import collections +import math +from dataclasses import asdict + +import pytest +import torch +from lightllm.common.ssm_types import SSMRunConfig + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor +from lightllm.common.attention_state.commit import StateAcceptance, 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 = ReplaySSMExecutor(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) + else: + ssm_backend = CompactSSMExecutor(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 + 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 SSMRunConfig(**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 == SSMRunConfig(8, 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, manager._ssm, 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") + 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)] + acceptance = StateAcceptance( + rows[starts.long()], + torch.tensor(counts, device="cuda"), + tuple(StateCommitBatch(len(b[-1]), b[2].state_commits) for b in buckets), + ) + manager.commit_bounded_att_state(acceptance) + 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..756774c8b9 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,70 @@ +import pytest +import torch +from lightllm.common.ssm_types import SSMRunConfig + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor + + +@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 = ReplaySSMExecutor(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 = ReplaySSMExecutor(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 = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + for tensor in vars(ssm_backend).values(): + if isinstance(tensor, torch.Tensor): + tensor.zero_() + tokens = 3 * width + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + bias = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width 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 == SSMRunConfig(32, 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..5d605dae15 --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,231 @@ +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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("destination", ["cpu", "cuda"]) +def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dtype, destination): + state, _, _ = make_state(cls, mode, 2, dtype) + state.initialize_page_pools(3, device=destination) + state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) + if mode == "replay": + state._ssm.raw_keys.zero_() + state._ssm.raw_values.zero_() + state._ssm.gates.fill_(-0.5) + state._ssm.betas.zero_() + state._ssm.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + requests = [2, 0, 1] + expected = [(s.conv.clone(), s.ssm.clone()) for s in map(state.get_accepted_state, requests)] + before = state.ssm_state.clone() + cursors = state._ssm.cursors.clone() if mode == "replay" else None + + def no_snapshot(*args): + raise AssertionError("batch export must write directly to the destination") + + monkeypatch.setattr(state, "get_accepted_state", no_snapshot) + monkeypatch.setattr(state._ssm, "get_accepted_state", no_snapshot) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.save_bounded_att_runtime_to_pages(requests, [8] * 3, "big", [1, 2, 0]) + stream.synchronize() + pool = state.big_pages_pool + for slot, (conv, ssm) in zip([1, 2, 0], expected): + torch.testing.assert_close(pool.conv_state[slot].cpu(), conv.cpu(), atol=0, rtol=0) + torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state._ssm.cursors, cursors) + target, _, _ = make_state(cls, "native", 2, dtype) + target.mtp_state_index.fill_(2) + target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) + for req, (conv, ssm) in enumerate(expected): + actual = target.get_accepted_state(req) + torch.testing.assert_close(actual.conv, conv, atol=0, rtol=0) + torch.testing.assert_close(actual.ssm, ssm, atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state._ssm.cursors, cursors) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_native_checkpoint_scratch_is_bounded(runtime_start_args): + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1, 2, 3), + global_k_heads=1, + global_v_heads=2, + key_dim=128, + value_dim=128, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = GDNState(config, 65, mtp_step=2, device="cuda") + state.initialize_page_pools(64) + requests = list(range(64)) + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + assert torch.cuda.max_memory_allocated() - before < 1024 * 1024 + + +def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): + # Exercise layout and accepted-state selection independently of the CUDA copy kernel. + monkeypatch.setattr( + import_module("lightllm.common.attention_state.bounded_att.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..a546814438 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,338 @@ +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 StateAcceptance, 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 ReplaySSMExecutor + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor + + ssm_backend = object.__new__(CompactSSMExecutor if mode == "compact" else ReplaySSMExecutor) + 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_acceptance(batches=()): + return StateAcceptance(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), 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( + owner._ssm, + "accept_updates", + lambda reqs, accepted, run_config: calls.append((reqs.tolist(), accepted, run_config)), + ) + acceptance = make_acceptance( + (StateCommitBatch(1, (commits[0],)), StateCommitBatch(0), StateCommitBatch(2, (commits[1],))) + ) + owner.commit_bounded_att_state(acceptance) + assert owner.mtp_state_index.tolist() == [1, 2, 0, 0] + 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_batches"]) +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),) + acceptance = make_acceptance(() if failure == "no_batches" else (StateCommitBatch(3, commits),)) + with pytest.raises(ValueError, match="compact"): + owner.commit_bounded_att_state(acceptance) + 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)) + acceptance = make_acceptance((StateCommitBatch(3, (commit,)),)) + with pytest.raises(ValueError, match="unknown state"): + acceptance.validate({"draft0"}) + with pytest.raises(ValueError, match="layout"): + StateAcceptance(torch.tensor([1, 2]), torch.tensor([1])).validate() + with pytest.raises(ValueError, match="partition"): + make_acceptance((StateCommitBatch(2),)).validate({"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(StateAcceptance(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, "commit_ssm_prefix", lambda *args: calls.append(args)) + monkeypatch.setattr(owner._ssm, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(make_acceptance()) + assert len(calls) == 1 + assert calls[0][3] is owner._ssm.cursors + assert calls[0][0] is owner.mtp_state_index + + +@pytest.mark.parametrize("change", ["owner", "executor", "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 == "executor": + source.executor = copy(owner._ssm) + 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 + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_acceptance_keeps_request_order_with_strided_inputs(device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(mode="replay", step=3, device=device) + owner.mtp_state_index.fill_(2) + owner._ssm.cursors.fill_(8) + reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] + lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] + owner.commit_bounded_att_state(StateAcceptance(reqs, lengths)) + assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] + assert owner._ssm.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 1da9c1e286..622c8087ad 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -22,6 +22,19 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv(key, value) +@pytest.fixture +def runtime_start_args(monkeypatch: pytest.MonkeyPatch): + import json + from dataclasses import asdict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 530b95436a..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_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 8a1d70e3ac..60f94ed27c 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -125,17 +125,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)] @@ -143,10 +144,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -184,7 +185,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -203,11 +204,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) @@ -223,7 +232,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, @@ -237,6 +246,26 @@ 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 dynamic: - assert torch.equal(conv[3], old_conv[3]) + if mode == "native": + checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] + torch.testing.assert_close( + selected.ssm_state[0, :checked_slots], expected_ssm[:checked_slots], atol=2e-3, rtol=1e-2 + ) + else: + 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 and mtp_step: + assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index e72e4d6bee..0258b93a6e 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..72b2f25f51 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,33 @@ 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"} 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..d4362cf09e --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,219 @@ +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", "bfloat16", 2, "sensenova_loop", 16, None), + ("replay", "float32", 0, "sensenova_loop", 16, None), + ("replay", "bfloat16", 2, "sensenova_loop", 8, 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) + + +@pytest.mark.parametrize("model", ["glm5_next", "glm5_next_text"]) +def test_kda_rejects_parallel_at_startup(monkeypatch, model): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"compute_mode": "parallel"}, + mtp_step=2, + linear_att_ssm_data_type="bfloat16", + ) + with pytest.raises(AssertionError, match="requires GDN, not KDA"): + _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]) From 398cc21e4740a0a6f77c57ed478fbfb2f812cf12 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 8 Nov 2026 19:30:57 +0800 Subject: [PATCH 02/54] refactor: remove retired ReplaySSM modes --- .../attention_state/bounded_att/__init__.py | 6 +- .../attention_state/bounded_att/base_state.py | 4 - .../bounded_att/conv_ssm_state.py | 20 +- .../bounded_att/gdn/__init__.py | 3 +- .../bounded_att/gdn/create_gdn_state.py | 3 +- .../bounded_att/gdn/gdn_state.py | 8 +- .../attention_state/bounded_att/interface.py | 4 +- .../bounded_att/kda/__init__.py | 3 +- .../bounded_att/kda/create_kda_state.py | 3 +- .../bounded_att/kda/kda_state.py | 8 +- .../sliding_window/fp8_e8m0_swa_state.py | 4 +- .../bounded_att/ssm/__init__.py | 16 +- .../attention_state/bounded_att/ssm/base.py | 2 +- .../bounded_att/ssm/compact.py | 192 -------- .../attention_state/bounded_att/ssm/native.py | 2 +- .../attention_state/bounded_att/ssm/replay.py | 45 +- .../window_compressor_state.py | 4 +- lightllm/common/attention_state/manager.py | 2 + .../triton_kernel/linear_att/compact_ssm.py | 112 ----- .../triton_kernel/linear_att/replayssm.py | 2 - .../linear_att/replayssm_common.py | 3 +- .../linear_att/replayssm_compute.py | 67 +-- .../linear_att/replayssm_fold.py | 56 +-- .../linear_att/replayssm_materialize.py | 30 +- .../triton_kernel/linear_att/ssm_autotune.py | 71 ++- lightllm/server/api_cli.py | 10 +- lightllm/server/api_start.py | 24 +- .../server/router/model_infer/infer_batch.py | 2 + test/benchmark/kernels/benchmark_kda_ssm.py | 7 +- .../test_hybrid_layer_mapping.py | 2 +- .../basemodel/test_cuda_graph_autotune.py | 70 --- .../linear_att/test_replayssm.py | 175 ++------ .../linear_att/test_replayssm_grouped_keys.py | 81 ---- .../linear_att/test_replayssm_parallel.py | 172 ------- .../linear_att/test_ssm_autotune.py | 423 ------------------ .../test_hybrid_req_manager_pd_sharing.py | 2 +- unit_tests/common/test_linear_state.py | 4 +- .../common/test_linear_state_contract.py | 109 +---- unit_tests/models/glm5_next/test_cache.py | 4 +- unit_tests/models/glm5_next/test_mtp.py | 7 +- unit_tests/models/glm5_next/test_pd_cache.py | 4 +- unit_tests/server/test_hybrid_start_args.py | 8 +- unit_tests/server/test_ssm_state_mode.py | 64 +-- 43 files changed, 224 insertions(+), 1614 deletions(-) delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/compact.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py index 9e5633805e..4efb909f0a 100644 --- a/lightllm/common/attention_state/bounded_att/__init__.py +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -2,9 +2,9 @@ from .kda.create_kda_state import create_kda_state from .conv_ssm_types import ConvSSMStateConfig from .gdn.gdn_pages_pool import GDNPagesPool -from .gdn.gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState +from .gdn.gdn_state import GDNNativeState, GDNReplayState, GDNState from .kda.kda_pages_pool import KDAPagesPool -from .kda.kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState +from .kda.kda_state import KDANativeState, KDAReplayState, KDAState from .base_state import BoundedAttentionState from .sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState @@ -17,13 +17,11 @@ "ConvSSMStateConfig", "GDNState", "GDNNativeState", - "GDNCompactState", "GDNReplayState", "GDNPagesPool", "create_gdn_state", "KDAState", "KDANativeState", - "KDACompactState", "KDAReplayState", "KDAPagesPool", "create_kda_state", diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index dca8f4e0f3..6a96e76dc9 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -7,7 +7,3 @@ class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" - - def validate_bounded_att_commit(self, acceptance): - - 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 index 614e2629de..dd2b5e1ace 100644 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -8,6 +8,7 @@ from lightllm.utils.log_utils import init_logger +from ..commit import StateAcceptance from .ssm import create_ssm_executor from .ssm.forward import SSMForward from .ssm.export import export_indexed_state @@ -204,6 +205,7 @@ def save_bounded_att_runtime_to_pages( state_token_lens: List[int], page_kind: PageKind, page_slot_ids: List[int], + flush_history: bool = False, ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance( @@ -211,8 +213,18 @@ def save_bounded_att_runtime_to_pages( ), "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) - self.export_accepted_state( - req_indices, [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] + # Hybrid radix pages are durable checkpoints. ReplaySSM can fold the + # accepted ring into the live canonical state and the page in one + # kernel, so the next decode starts from an empty history. MTP keeps + # a separate Conv candidate layout and retains the old behavior. + replay_flush = flush_history and self.mtp_state_index is None + targets = [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] + export_indexed_state(self.conv_state, req_indices, [target.conv for target in targets], self.mtp_state_index) + self._ssm.export_accepted_state( + req_indices, + [target.ssm for target in targets], + self.mtp_state_index, + flush=replay_flush, ) def restore_bounded_att_runtime_from_pages( @@ -236,7 +248,7 @@ def restore_bounded_att_runtime_from_pages( copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def validate_bounded_att_commit(self, acceptance): + def validate_bounded_att_commit(self, acceptance: StateAcceptance): assert ( acceptance.req_indices.device == self.ssm_state.device ), "commit inputs must match the state device; CUDA states require CUDA tensors" @@ -245,7 +257,7 @@ def validate_bounded_att_commit(self, acceptance): return validator(acceptance, self.state_name) if validator is not None else None return None - def commit_bounded_att_state(self, acceptance): + def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: commit_data = self.validate_bounded_att_commit(acceptance) if self.mtp_state_index is not None: self._ssm.commit_validated(acceptance, self.mtp_state_index, self.state_name, commit_data=commit_data) diff --git a/lightllm/common/attention_state/bounded_att/gdn/__init__.py b/lightllm/common/attention_state/bounded_att/gdn/__init__.py index 03e545bfa2..054e9bb519 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/__init__.py +++ b/lightllm/common/attention_state/bounded_att/gdn/__init__.py @@ -1,12 +1,11 @@ from .create_gdn_state import create_gdn_state from .gdn_pages_pool import GDNPagesPool -from .gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState +from .gdn_state import GDNNativeState, GDNReplayState, GDNState __all__ = [ "create_gdn_state", "GDNPagesPool", "GDNState", "GDNNativeState", - "GDNCompactState", "GDNReplayState", ] 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 cfa58d42e9..35eb5147e3 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .gdn_state import GDNCompactState, GDNNativeState, GDNReplayState, GDNState +from .gdn_state import GDNNativeState, GDNReplayState, GDNState def create_gdn_state( @@ -21,7 +21,6 @@ def create_gdn_state( """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" state_cls = { "native": GDNNativeState, - "compact": GDNCompactState, "replay": GDNReplayState, }.get(ssm_state_mode) if state_cls is None: 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 04c7a52bce..12cb426a7c 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -1,4 +1,4 @@ -"""GDN 运行态及其三种 SSM 执行策略。""" +"""GDN 运行态及其 SSM 执行策略。""" from ..conv_ssm_state import _ConvSSMState from .gdn_pages_pool import GDNPagesPool @@ -17,12 +17,6 @@ class GDNNativeState(GDNState): state_mode = "native" -class GDNCompactState(GDNState): - """GDN compact 接受前缀重建策略。""" - - state_mode = "compact" - - class GDNReplayState(GDNState): """GDN replay history 策略。""" diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index 97e6261787..983d8e78c5 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -76,11 +76,13 @@ def save_bounded_att_runtime_to_pages( state_token_lens: List[int], page_kind: PageKind, page_slot_ids: List[int], + flush_history: bool = False, ) -> None: """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 - 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 + ``flush_history`` 用于持久化 radix checkpoint：支持的 State 可在写页时 + 同步折叠临时 history；前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 """ ... diff --git a/lightllm/common/attention_state/bounded_att/kda/__init__.py b/lightllm/common/attention_state/bounded_att/kda/__init__.py index 3fcab8e5d1..498ba3b227 100644 --- a/lightllm/common/attention_state/bounded_att/kda/__init__.py +++ b/lightllm/common/attention_state/bounded_att/kda/__init__.py @@ -1,12 +1,11 @@ from .create_kda_state import create_kda_state from .kda_pages_pool import KDAPagesPool -from .kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState +from .kda_state import KDANativeState, KDAReplayState, KDAState __all__ = [ "create_kda_state", "KDAPagesPool", "KDAState", "KDANativeState", - "KDACompactState", "KDAReplayState", ] 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 99f7246ea2..ba4614da27 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .kda_state import KDACompactState, KDANativeState, KDAReplayState, KDAState +from .kda_state import KDANativeState, KDAReplayState, KDAState def create_kda_state( @@ -21,7 +21,6 @@ def create_kda_state( """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" state_cls = { "native": KDANativeState, - "compact": KDACompactState, "replay": KDAReplayState, }.get(ssm_state_mode) if state_cls is None: 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 d7562455f2..6723e8aeec 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -1,4 +1,4 @@ -"""KDA 运行态及其三种 SSM 执行策略。""" +"""KDA 运行态及其 SSM 执行策略。""" from ..conv_ssm_state import _ConvSSMState from .kda_pages_pool import KDAPagesPool @@ -17,12 +17,6 @@ class KDANativeState(KDAState): state_mode = "native" -class KDACompactState(KDAState): - """KDA compact 接受前缀重建策略。""" - - state_mode = "compact" - - class KDAReplayState(KDAState): """KDA replay history 策略。""" 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 f6aadf7b61..d1455cdb34 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 @@ -9,6 +9,7 @@ import torch +from ...commit import StateAcceptance from ...utils import deepcopy_result from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState @@ -153,6 +154,7 @@ def save_bounded_att_runtime_to_pages( state_token_lens: list[int], page_kind: PageKind, page_slot_ids: list[int], + flush_history: bool = False, ) -> None: pool = self._pool(page_kind=page_kind) for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): @@ -179,7 +181,7 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, acceptance) -> None: + def commit_bounded_att_state(self, acceptance: StateAcceptance) -> 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 index ab61fc5f7d..15dcd1501a 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/__init__.py +++ b/lightllm/common/attention_state/bounded_att/ssm/__init__.py @@ -2,14 +2,13 @@ def create_ssm_executor(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound): - """按统一 state layout 创建 native、compact 或 replay executor。 + """按统一 state layout 创建 native 或 replay executor。 - 三种模式共享 canonical SSM state；差异只在 decode 中间量和接受阶段的 - state 更新方式。native 为每个 MTP 位置分配物理槽位，compact/replay - 则为每个请求分配一个槽位并在各自 executor 内保存临时数据。 + native 为每个 MTP 位置分配物理槽位，replay 为每个请求分配一个槽位， + 并在 executor 内保存短 history。 """ - if mode not in ("native", "compact", "replay"): + if mode not in ("native", "replay"): raise ValueError(f"unknown SSM state mode: {mode}") slots = mtp_step + 1 if mode == "native" else 1 state = torch.zeros( @@ -23,10 +22,6 @@ def create_ssm_executor(config, request_capacity, mtp_step, device, mode, replay return NativeSSMExecutor(state, mtp_step + 1, **kwargs) kwargs.update(num_key_heads=config.num_k_heads) - if mode == "compact": - from .compact import CompactSSMExecutor - - return CompactSSMExecutor(state, mtp_step + 1, config.conv_dtype, **kwargs) from .replay import ReplaySSMExecutor replay_config = replay_config or {} @@ -35,7 +30,6 @@ def create_ssm_executor(config, request_capacity, mtp_step, device, mode, replay 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"), + compute_mode=replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline"), **kwargs, ) diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py index 984a776d26..3070856c7a 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/base.py +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -63,7 +63,7 @@ def get_accepted_state(self, req_idx, accepted_index=None): return self.state[:, req_idx] - def export_accepted_state(self, req_indices, targets, accepted_indices=None): + def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): """把 accepted state 写入调用方提供的目标页，不负责分配目标。""" export_indexed_state(self.state, req_indices, targets) diff --git a/lightllm/common/attention_state/bounded_att/ssm/compact.py b/lightllm/common/attention_state/bounded_att/ssm/compact.py deleted file mode 100644 index 1c1b1d3d46..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/compact.py +++ /dev/null @@ -1,192 +0,0 @@ -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 SSMExecutor, update_accepted_indices - - -class CompactSSMExecutor(SSMExecutor): - """保存 verify 中间量，并在接受阶段按原 forward 配置重建 state。""" - - 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 = run_config or SSMRunConfig(8, 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: - """选择与 verify layout 一致的 compact kernel 配置。""" - - config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) - return 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, - ): - """执行 compact verify；该模式只接受带 ``cu_seqlens`` 的 speculative 输入。""" - - 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) - _, 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): - """按 accepted index 重建请求前缀，并使用 verify 时的配置保持舍入一致。""" - 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 - 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): - """把本次 verify 的配置作为 compact 接受提交 payload。""" - - return run_config - - def validate_commit(self, acceptance, state_name): - """确认接受结果为每个 microbatch 提供了 compact forward 配置。""" - - if not acceptance.batches: - raise ValueError("compact acceptance requires forward commit batches") - partitions = tuple(acceptance.partitions(state_name)) - for _, config in partitions: - if not isinstance(config, SSMRunConfig): - raise ValueError(f"missing compact forward configuration for {state_name}") - return partitions - - def commit(self, acceptance, indices, state_name): - """先推进通用接受索引，再逐 microbatch 重建 compact state。""" - - commit_data = self.validate_commit(acceptance, state_name) - self.commit_validated(acceptance, indices, state_name, commit_data=commit_data) - - def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): - """使用已验证的 forward 配置完成 compact state 重建。""" - - if commit_data is None: - commit_data = self.validate_commit(acceptance, state_name) - update_accepted_indices(acceptance, indices) - for reqs, config in commit_data: - self.accept_updates(reqs, indices, run_config=config) diff --git a/lightllm/common/attention_state/bounded_att/ssm/native.py b/lightllm/common/attention_state/bounded_att/ssm/native.py index 9303efcc2e..92ee179b04 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/native.py +++ b/lightllm/common/attention_state/bounded_att/ssm/native.py @@ -29,7 +29,7 @@ def get_accepted_state(self, req_idx, accepted_index=None): index = req_idx * self.slots_per_request + accepted_index return self.state.index_select(1, index.long().view(1))[:, 0] - def export_accepted_state(self, req_indices, targets, accepted_indices=None): + def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): """把 native state 按请求和接受位置导出到调用方目标页。""" export_indexed_state( diff --git a/lightllm/common/attention_state/bounded_att/ssm/replay.py b/lightllm/common/attention_state/bounded_att/ssm/replay.py index 629bd35da8..3378fb3c04 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/replay.py +++ b/lightllm/common/attention_state/bounded_att/ssm/replay.py @@ -35,26 +35,16 @@ def __init__( gate_lower_bound=-5.0, compute_mode="inline", run_config=None, - fold_method="raw", ): """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行 - 方式，``fold_method`` 选择 raw 或 derived history 的折叠路径。 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 """ 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 + assert compute_mode in ("inline", "precompute") self.state = state self.kda = kda self.gate_lower_bound = gate_lower_bound @@ -76,9 +66,8 @@ def __init__( 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. + # BF16 can use a narrower short-history buffer for the precompute path; + # FP32 remains the safe default for all other shapes. history_dtype = ( activation_dtype if activation_dtype == torch.bfloat16 @@ -125,7 +114,6 @@ def prepare_decode(self, reqs, cu_seqlens=None): 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) @@ -153,10 +141,6 @@ def prepare_decode(self, reqs, cu_seqlens=None): 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, ) @@ -185,9 +169,11 @@ def merge_accepted_updates(self, reqs): 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, output_strides=None): + def _materialize_accepted_state(self, reqs, output, snapshot, output_strides=None, flush_export=False): """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.state.shape config = self.run_config or SSMRunConfig(32, 4) bv = config.BV @@ -211,6 +197,7 @@ def _materialize_accepted_state(self, reqs, output, snapshot, output_strides=Non bv, snapshot, self.kda, + FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, num_warps=config.num_warps, num_stages=config.num_stages, @@ -224,7 +211,7 @@ def get_accepted_state(self, req_idx, accepted_index=None): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def export_accepted_state(self, req_indices, targets, accepted_indices=None): + def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" if not req_indices: @@ -236,18 +223,15 @@ def export_accepted_state(self, req_indices, targets, accepted_indices=None): self._materialize_accepted_state( upload_indices(req_indices, self.state.device, torch.int32), upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=True, + snapshot=not flush, + flush_export=flush, output_strides=strides, ) def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" - if self.compute_mode == "parallel": - config = self.run_config or SSMRunConfig(128, 2, 1) - else: - config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - return config + return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) def forward( self, @@ -274,9 +258,7 @@ def forward( 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" + precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode == "precompute" # 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) @@ -318,7 +300,6 @@ def forward( bw, 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, 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 958295b1ba..54d92a2e4e 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 @@ -12,6 +12,7 @@ from lightllm.utils.envs_utils import get_env_start_args +from ...commit import StateAcceptance from ...utils import deepcopy_result from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState @@ -173,6 +174,7 @@ def save_bounded_att_runtime_to_pages( state_token_lens: list[int], page_kind: PageKind, page_slot_ids: list[int], + flush_history: bool = False, ) -> None: """边界必须完成一个产出步长；仅保留下一步仍需读取的窗口重叠部分。""" assert all( @@ -210,7 +212,7 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, acceptance) -> None: + def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 4125ff727e..5bb052374a 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -346,6 +346,7 @@ def save_bounded_att_runtime_to_pages( state_token_lens: List[int], page_kind: PageKind, page_slot_ids: List[int], + flush_history: bool = False, ) -> None: """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" bounded_states = self.get_bounded_attention_states() @@ -364,6 +365,7 @@ def save_bounded_att_runtime_to_pages( state_token_lens=state_token_lens, page_kind=page_kind, page_slot_ids=page_slot_ids, + flush_history=flush_history, ) def restore_bounded_att_runtime_from_pages( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py deleted file mode 100644 index c16a9d7983..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py +++ /dev/null @@ -1,112 +0,0 @@ -# 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/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 876cfb9101..81b2804132 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -11,7 +11,6 @@ _replay, _replay_gdn_window, _replay_sequential, - _solve_gdn_deltas, ) from .replayssm_fold import _accept_updates, _fold_history, _prepare_fold, _use_separate_fold from .replayssm_materialize import _materialize @@ -26,6 +25,5 @@ "_replay_gdn_window", "_replay_sequential", "_replay_update", - "_solve_gdn_deltas", "_use_separate_fold", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py index a54b264947..84c1af9abe 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py @@ -7,8 +7,7 @@ @triton.jit def _replay_update(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): - # 所有 replay fold 路径共用这一条递推，保证 raw/derived/history - # materialize 使用相同的数值顺序。 + # Replay folding and materialization use the same update order. raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) decay = tl.exp(g) if KDA: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py index de216db8a7..604c09f1e0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py @@ -25,7 +25,6 @@ def _project_replay_queries( BW: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, - PARALLEL_SOLVE: tl.constexpr, BQK: tl.constexpr, base, end, @@ -39,7 +38,7 @@ def _project_replay_queries( start, state, ): - # 预计算 verify window 的 q/k 投影；仅 precompute/parallel 路径调用。 + # 预计算 verify window 的 q/k 投影。 tt = tl.arange(0, BQK) tokens = start + tt % BW valid = (tt < 2 * BW) & (tokens < end) @@ -55,10 +54,7 @@ def _project_replay_queries( 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.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 @@ -94,21 +90,6 @@ def _project_replay_queries( return qk, projections, tt, tokens, valid, h -@triton.jit -def _solve_gdn_deltas(power, rhs, BW: tl.constexpr, projection_dtype: tl.constexpr, projection_precision: tl.constexpr): - # 解同一 verify window 内的严格下三角 delta 耦合，不读写 history。 - wi = tl.arange(0, power.shape[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") - return tl.dot(inverse.to(projection_dtype), rhs.to(projection_dtype), input_precision=projection_precision) - - @triton.jit def _replay_gdn_window( Kp, @@ -135,8 +116,6 @@ def _replay_gdn_window( L: tl.constexpr, BV: tl.constexpr, BW: tl.constexpr, - PARALLEL_SOLVE: tl.constexpr, - BQK: tl.constexpr, BR: tl.constexpr, base, end, @@ -199,32 +178,14 @@ def _replay_gdn_window( 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) - solved = _solve_gdn_deltas(power, rhs, BW, projection_dtype, 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) + 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 ) @@ -434,11 +395,10 @@ def _replay( BW: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, - PARALLEL_SOLVE: tl.constexpr, BQK: tl.constexpr, BR: tl.constexpr, ): - """执行一个请求窗口的 replay decode，并写入 raw/derived history。 + """执行一个请求窗口的 replay decode，并写入 history。 当 history 即将跨半区时，先按 raw input 折叠到 canonical state；随后 根据 compute mode 选择预计算投影或逐 token 递推。请求数属于运行时布局， @@ -468,7 +428,7 @@ def _replay( 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 not PARALLEL_SOLVE and n + end - start > L: + if 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) @@ -505,7 +465,6 @@ def _replay( BW, KDA, LOWER, - PARALLEL_SOLVE, BQK, base, end, @@ -548,8 +507,6 @@ def _replay( L, BV, BW, - PARALLEL_SOLVE, - BQK, BR, base, end, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 92a1108b0c..3281c5c402 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -27,7 +27,7 @@ def _accept_updates( @triton.jit -def _use_separate_fold(valid, length, count, N, L: tl.constexpr, WIDTH: tl.constexpr, FORCE: tl.constexpr): +def _use_separate_fold(valid, length, count, N, L: tl.constexpr, WIDTH: tl.constexpr): # 只决定本轮是否单独启动 fold；真正的 history 合并由 _fold_history 完成。 small_batch: tl.constexpr = 12 dense_batch: tl.constexpr = 256 @@ -37,8 +37,6 @@ def _use_separate_fold(valid, length, count, N, L: tl.constexpr, WIDTH: tl.const if N >= dense_batch: short = tl.sum((valid & (length > 0) & (length < WIDTH)).to(tl.int32), axis=0) split = split & ((live < dense_batch) | (short > 0)) - if FORCE: - split = count > 0 return split @@ -56,7 +54,6 @@ def _prepare_fold( VARLEN: tl.constexpr, BLOCK: tl.constexpr, WIDTH: tl.constexpr, - SEPARATE_FOLD: tl.constexpr, ): # 输出 Positions/ForwardCursors 供 decode 使用，Active/N 供后续 fold # 使用；只改写本轮请求对应的 cursor，不触碰 canonical state。 @@ -77,7 +74,7 @@ def _prepare_fold( 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) - split = _use_separate_fold(valid, length, count, N, L, WIDTH, SEPARATE_FOLD) + split = _use_separate_fold(valid, length, count, N, L, WIDTH) tl.store(ForwardCursors + req, tl.where(overflow & split, (cursor & (2 * L)) ^ (2 * L), cursor), valid) tl.store(Active + N, tl.where(split, count, 0)) @@ -107,10 +104,6 @@ def _fold_history( GS: tl.constexpr, BS: tl.constexpr, KDA: tl.constexpr, - Keys, - Deltas, - DERIVED: tl.constexpr, - BL: tl.constexpr, ): # 只消费 Active 中的 fold 请求，把 history 合并回 canonical state； # cursor 的 phase 已由 _prepare_fold 切换，写入后新 history 从同一半区继续。 @@ -134,40 +127,21 @@ def _fold_history( 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( + 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 ) - 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) + 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)) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py index 7f82d9f19e..3b19a13e6c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py @@ -26,10 +26,12 @@ def _materialize( BV: tl.constexpr, SNAPSHOT: tl.constexpr, KDA: tl.constexpr, + FLUSH_EXPORT: tl.constexpr = False, OUTPUT_STRIDES: tl.constexpr = None, ): # 将 cursor 指向的 accepted history 应用到 canonical state；SNAPSHOT - # 只写独立输出，merge 路径则原地更新请求槽位。 + # 只写独立输出，merge 路径则原地更新请求槽位。FLUSH_EXPORT 同时 + # 更新 canonical state 和缓存目标，并在同一个 kernel 中清空 cursor。 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: @@ -37,7 +39,7 @@ def _materialize( cursor = tl.load(Cursors + req) n = cursor % (2 * L) base = (cursor // (2 * L)) * L - if n == 0 and not SNAPSHOT: + if n == 0 and not SNAPSHOT and not FLUSH_EXPORT: return linear_layer_idx, head = lh // HV, lh % HV slot = (linear_layer_idx * SLOTS + req) * HV + head @@ -58,16 +60,30 @@ def _materialize( 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, :] - if OUTPUT_STRIDES is not None: + if FLUSH_EXPORT: target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) - out = ( + target_out = ( target + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + head * OUTPUT_STRIDES[1] + kk[:, None] * OUTPUT_STRIDES[2] + vv[None, :] * OUTPUT_STRIDES[3] ) + tl.store(target_out, state, (kk[:, None] < K) & (vv[None, :] < V)) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + if iv == 0 and lh == 0: + tl.store(Cursors + req, 0) else: - 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)) + if OUTPUT_STRIDES is not None: + target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) + out = ( + target + + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + + head * OUTPUT_STRIDES[1] + + kk[:, None] * OUTPUT_STRIDES[2] + + vv[None, :] * OUTPUT_STRIDES[3] + ) + else: + 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 index ac7e920aaa..bcfe691d42 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -31,23 +31,18 @@ def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens): "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, + "history": ssm_backend.capacity, "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" - ) + key["history_dtype"] = str(ssm_backend.keys.dtype) + 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 @@ -86,27 +81,16 @@ def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=Non ) _, _, 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, - ) + 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, + ) # Benchmark callbacks must not recursively select a configuration. scratch._config_is_fixed = True inputs = [ @@ -122,8 +106,6 @@ def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=Non 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, @@ -135,7 +117,7 @@ def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=Non a_log=inputs[5], bias=inputs[6], cu_seqlens=cu if cu_seqlens is not None else None, - workload=(reqs, accepted, partial), + workload=(reqs, accepted), ) @@ -154,19 +136,16 @@ def select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None 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. + reqs, accepted = workload + 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 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) + ssm_backend.accept_updates(reqs, accepted) return run_config diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 6911939f9e..4965e91824 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -987,17 +987,17 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: ) parser.add_argument( "--ssm_state_mode", - choices=["native", "compact", "replay"], + choices=["native", "replay"], default="native", - help="SSM state storage: native snapshots, compact MTP records, or deferred GDN/KDA replay (FP32/BF16).", + help="SSM state storage: native snapshots or deferred GDN/KDA ReplaySSM (FP32/BF16).", ) parser.add_argument( "--ssm_replay_config", type=json.loads, default={}, - help='Replay-only JSON object: {"cache_len":16,"compute_mode":"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.", + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"precompute"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute. " + "MTP defaults to precompute, single-token decode defaults to inline.", ) parser.add_argument( "--linear_att_ssm_data_type", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 302c04d4de..95379e5447 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -44,12 +44,12 @@ def _set_envs_and_config(args: StartArgs): def _validate_and_normalize_ssm_config(args: StartArgs): - assert args.ssm_state_mode in ("native", "compact", "replay"), "Invalid ssm_state_mode" + assert args.ssm_state_mode in ("native", "replay"), "Invalid ssm_state_mode" config = args.ssm_replay_config assert isinstance(config, dict), "ssm_replay_config must be a JSON object" assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" if args.ssm_state_mode == "replay": - defaults = {"cache_len": 16, "compute_mode": "inline", "fold_method": "raw"} + defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} unknown = config.keys() - defaults.keys() assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" config = {**defaults, **config} @@ -60,14 +60,7 @@ def _validate_and_normalize_ssm_config(args: StartArgs): 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" + assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" args.ssm_replay_config = config logger.info(f"SSM mode: replay, config: {config}") else: @@ -96,14 +89,9 @@ def _validate_and_normalize_ssm_config(args: StartArgs): "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" - if config["compute_mode"] == "parallel": - assert model_type not in ("glm5_next", "glm5_next_text"), "Parallel ReplaySSM requires GDN, not KDA" + 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): diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 942bca42cd..ae79822bf1 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -433,6 +433,7 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR state_token_lens=[req.get_chuncked_input_token_len() for req, _ in selected], page_kind="big", page_slot_ids=[slot for _, slot in selected], + flush_history=True, ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -454,6 +455,7 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR state_token_lens=[req.hybrid_cache_len], page_kind="small", page_slot_ids=[dst_buffer_idx], + flush_history=True, ) return diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 3b702f206d..62088a6d91 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -16,7 +16,6 @@ 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 ReplaySSMExecutor -from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs @@ -37,9 +36,7 @@ def workload(args, batch, width, dtype, mode, acceptance): 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 = CompactSSMExecutor(state, width, q.dtype, kda=True) - elif mode != "native": + if mode != "native": capacity = max(8, triton.next_power_of_2(width)) ssm_backend = ReplaySSMExecutor(state, capacity, width, q.dtype, kda=True, compute_mode=mode) if ssm_backend is not None: @@ -112,7 +109,7 @@ def main(): 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"]: + for mode in ["native", "inline", "precompute"] if width > 1 else ["native", "inline"]: ssm_backend, cycle, size = workload( args, batch, width, getattr(torch, dtype_name), mode, acceptance ) 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 4580453297..86c5aa9282 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -9,7 +9,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs -@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("draft_layers", [0, 1, 3]) @pytest.mark.parametrize( "layer_types,expected_indices,linear_layers", diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index d417381753..3944c0e229 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -129,73 +129,3 @@ 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 CompactSSMExecutor - 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_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index ee83837a04..e88c021fbf 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -59,7 +59,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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. + # an overflowed history 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] @@ -81,6 +81,37 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state torch.testing.assert_close(ssm_backend.get_accepted_state(0)[0].float(), expected_state, atol=128, rtol=0.01) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_flush_export_updates_checkpoint_and_clears_history(dtype): + state = torch.randn(1, 2, 1, 32, 64, device="cuda", dtype=dtype) * 0.01 + backend = ReplaySSMExecutor( + state, + capacity=4, + verify_width=1, + activation_dtype=torch.bfloat16, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + backend.raw_keys.normal_(0, 0.1) + backend.raw_values.normal_(0, 0.1) + backend.gates.fill_(-0.1) + backend.betas.fill_(0.5) + backend.cursors[0] = 2 + expected = backend.get_accepted_state(0).clone() + target = torch.empty_like(expected) + + backend.export_accepted_state([0], [target], flush=True) + torch.cuda.synchronize() + + assert backend.cursors[0].item() == 0 + torch.testing.assert_close(target, expected, atol=0, rtol=0) + torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) + + empty_history_target = torch.empty_like(expected) + backend.export_accepted_state([0], [empty_history_target], flush=True) + torch.cuda.synchronize() + torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) + + @pytest.mark.parametrize("capacity", [16, 32, 64]) @pytest.mark.parametrize("width", [1, 3, 5, 16]) @pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) @@ -437,114 +468,10 @@ def test_snapshot_without_history_copies_active_state(): 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 CompactSSMExecutor - - 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 = CompactSSMExecutor(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 CompactSSMExecutor - 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 = CompactSSMExecutor(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) +def test_autotune_candidates_match_recurrence(dtype, run_config): + test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) @pytest.mark.parametrize("save_big", [False, True]) @@ -638,28 +565,20 @@ def test_kda_replay_reference_fold_and_graph( ) -@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, - ) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_kda_tuning_candidates_match_reference(config, compute_mode): + test_replay_acceptance_flush_and_graph( + 3, + (128, 128), + 4, + torch.bfloat16, + run_config=config, + compute_mode=compute_mode, + head_geometry=(2, 2), + kda=True, + gate_shift=-3, + ) @pytest.mark.parametrize("kda", [False, True]) 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 deleted file mode 100644 index 04954d4361..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py +++ /dev/null @@ -1,81 +0,0 @@ -"""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 ReplaySSMExecutor -from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor - -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 ReplaySSMExecutor(state.clone(), 16, width, num_key_heads=key_heads) - return CompactSSMExecutor(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 deleted file mode 100644 index a423697c51..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py +++ /dev/null @@ -1,172 +0,0 @@ -import pytest -import torch - -from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor - -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 = ReplaySSMExecutor(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"): - ReplaySSMExecutor(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): - ReplaySSMExecutor(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 = ReplaySSMExecutor( - 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 deleted file mode 100644 index 4ffbcfc937..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ /dev/null @@ -1,423 +0,0 @@ -import collections -import math -from dataclasses import asdict - -import pytest -import torch -from lightllm.common.ssm_types import SSMRunConfig - -from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor -from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor -from lightllm.common.attention_state.commit import StateAcceptance, 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 = ReplaySSMExecutor(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) - else: - ssm_backend = CompactSSMExecutor(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 - 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 SSMRunConfig(**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 == SSMRunConfig(8, 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, manager._ssm, 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") - 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)] - acceptance = StateAcceptance( - rows[starts.long()], - torch.tensor(counts, device="cuda"), - tuple(StateCommitBatch(len(b[-1]), b[2].state_commits) for b in buckets), - ) - manager.commit_bounded_att_state(acceptance) - 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/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index cd0006dd45..f4399044e3 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -71,7 +71,7 @@ def pd_receiver(): @pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): state, resources, _ = make_state(cls=cls, mode=mode, device="meta") state.big_pages_pool, state.small_pages_pool = object(), object() diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 5d605dae15..e1adea50e9 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -29,7 +29,7 @@ def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device= @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("mode,step", [("native", 0), ("native", 2), ("replay", 0), ("replay", 2)]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( runtime_start_args, cls, mode, step, dtype @@ -94,7 +94,7 @@ def copy_on_cpu(*, copies): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("destination", ["cpu", "cuda"]) def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dtype, destination): diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index a546814438..b519b78519 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -1,6 +1,5 @@ import pickle from dataclasses import FrozenInstanceError -from types import SimpleNamespace import pytest import torch @@ -85,22 +84,19 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): 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") +def test_kda_backend_receives_gate_bound_at_construction(linear_config): + owner = KDAState(linear_config, 3, mtp_step=2, ssm_state_mode="replay", 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): +def test_replay_backend_publishes_detached_run_config(): from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor - from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMExecutor - ssm_backend = object.__new__(CompactSSMExecutor if mode == "compact" else ReplaySSMExecutor) + ssm_backend = object.__new__(ReplaySSMExecutor) 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" + ssm_backend.compute_mode = "inline" config = ssm_backend.get_run_config(*([None] * 7)) assert isinstance(config, SSMRunConfig) tuning["BV"] = 128 @@ -110,49 +106,15 @@ def test_ssm_backend_publishes_detached_run_config(mode): config.BV = 128 -def make_acceptance(batches=()): - return StateAcceptance(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), 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( - owner._ssm, - "accept_updates", - lambda reqs, accepted, run_config: calls.append((reqs.tolist(), accepted, run_config)), - ) - acceptance = make_acceptance( - (StateCommitBatch(1, (commits[0],)), StateCommitBatch(0), StateCommitBatch(2, (commits[1],))) - ) - owner.commit_bounded_att_state(acceptance) - assert owner.mtp_state_index.tolist() == [1, 2, 0, 0] - 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("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) +def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): + owner = cls(linear_config, request_capacity=2, mtp_step=mtp_step, ssm_state_mode="replay", device="cpu") + assert owner._ssm.compute_mode == expected -@pytest.mark.parametrize("failure", ["missing", "wrong_owner", "mutable_config", "no_batches"]) -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),) - acceptance = make_acceptance(() if failure == "no_batches" else (StateCommitBatch(3, commits),)) - with pytest.raises(ValueError, match="compact"): - owner.commit_bounded_att_state(acceptance) - assert owner.mtp_state_index.tolist() == [2] * 4 +def make_acceptance(batches=()): + return StateAcceptance(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), batches) def test_commit_partitions_and_resource_identity_are_validated(): @@ -229,53 +191,6 @@ def test_graph_layout_mismatch_rejected_before_copy(change): 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): diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index fa39331c26..0979efa560 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,7 +48,7 @@ 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"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) def test_hybrid_checkpoint_restore_and_packed_kv_copy( monkeypatch, small_page, @@ -56,8 +56,6 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( 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) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 60f94ed27c..4eafaad5d4 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -125,7 +125,7 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("role", ["main", "draft0"]) -@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("compact", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("replay", 2), ("native", 0), ("replay", 0)]) def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace @@ -256,10 +256,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod 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) + selected._ssm.accept_updates(ids, accepted) for req, length in zip(requests, lengths): torch.testing.assert_close( selected._ssm.get_accepted_state(req)[0], diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 0258b93a6e..ed219d2ede 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -114,7 +114,7 @@ def create(tp, target_mtp_step=None, mode="native"): @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"]) +@pytest.mark.parametrize("dest_mode", ["native", "replay"]) def test_pd_kv_and_runtime_state_roundtrip( make_mems, prefill_tp, @@ -124,8 +124,6 @@ def test_pd_kv_and_runtime_state_roundtrip( source_mode, dest_mode, ): - 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 diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 72b2f25f51..1c391c09cd 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -104,7 +104,7 @@ def test_whole_image_blocks_fit_batch_at_startup( _launch_subprocesses(args) -@pytest.mark.parametrize("mode,step", [("native", 0), ("compact", 4), ("replay", 0)]) +@pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): from pathlib import Path from lightllm.utils.petrel_helper import PetrelHelper @@ -131,4 +131,8 @@ def download_json(remote, local): _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"} + expected_compute_mode = "precompute" if step > 0 else "inline" + assert args.ssm_replay_config == { + "cache_len": 16, + "compute_mode": expected_compute_mode, + } diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index d4362cf09e..e60e71028d 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -12,27 +12,16 @@ [ ("native", "bfloat16", 0, "llama", 16, None), ("native", "bfloat16", 32, "llama", 16, None), - ("compact", "bfloat16", 2, "qwen3_5", 16, None), - ("compact", "bfloat16", 2, "sensenova_loop", 16, None), ("replay", "float32", 0, "sensenova_loop", 16, None), ("replay", "bfloat16", 2, "sensenova_loop", 8, 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"), ], ) @@ -58,7 +47,7 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er @pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) -@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("step", [15, 16]) def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) @@ -71,7 +60,7 @@ def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): @pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute", "parallel"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("cache_len", [32, 64]) def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) @@ -118,48 +107,6 @@ def test_replay_compute_mode_requires_replay(monkeypatch): 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) - - -@pytest.mark.parametrize("model", ["glm5_next", "glm5_next_text"]) -def test_kda_rejects_parallel_at_startup(monkeypatch, model): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) - args = StartArgs( - ssm_state_mode="replay", - ssm_replay_config={"compute_mode": "parallel"}, - mtp_step=2, - linear_att_ssm_data_type="bfloat16", - ) - with pytest.raises(AssertionError, match="requires GDN, not KDA"): - _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( @@ -173,7 +120,7 @@ def test_derived_replay_requires_parallel_configuration(): ({"cache_len": "16"}, "capacity"), ({"cache_len": 128}, "capacity"), ({"compute_mode": "auto"}, "compute_mode"), - ({"fold_method": "auto"}, "fold_method"), + ({"fold_method": "raw"}, "Unknown"), ], ) def test_invalid_replay_config(config, error): @@ -182,7 +129,7 @@ def test_invalid_replay_config(config, error): _validate_and_normalize_ssm_config(args) -@pytest.mark.parametrize("mode", ["native", "compact"]) +@pytest.mark.parametrize("mode", ["native"]) def test_replay_config_rejects_other_modes(mode): args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): @@ -197,7 +144,7 @@ def test_partial_replay_config_defaults_and_serialization(monkeypatch): ["--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"} + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute"} original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) restored = StartArgs(**json.loads(json.dumps(asdict(original)))) _validate_and_normalize_ssm_config(restored) @@ -210,7 +157,6 @@ def test_partial_replay_config_defaults_and_serialization(monkeypatch): [ ("--replayssm_cache_len", "16"), ("--replayssm_projection_mode", "inline"), - ("--replayssm_fold_method", "raw"), ("--ssm_replay_config", "{invalid}"), ], ) From ead7d04ec4385179bc946522a3421175469565cf Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 8 Nov 2026 19:40:02 +0800 Subject: [PATCH 03/54] refactor: name the SSM executor explicitly --- .../bounded_att/conv_ssm_state.py | 30 +++++++++------- .../attention_state/bounded_att/ssm/base.py | 4 +-- .../test_hybrid_layer_mapping.py | 4 +-- .../linear_att/test_acceptance_fusion.py | 24 ++++++------- .../linear_att/test_replayssm.py | 2 +- .../linear_att/test_replayssm_execution.py | 4 +-- .../test_hybrid_req_manager_pd_sharing.py | 20 +++++------ unit_tests/common/test_linear_state.py | 34 +++++++++---------- .../common/test_linear_state_contract.py | 14 ++++---- unit_tests/models/glm5_next/test_cache.py | 16 ++++----- unit_tests/models/glm5_next/test_mtp.py | 4 +-- unit_tests/models/glm5_next/test_pd_cache.py | 6 ++-- 12 files changed, 83 insertions(+), 79 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py index dd2b5e1ace..3726748512 100644 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -63,7 +63,7 @@ def __init__( dtype=config.conv_dtype, device=device, ) - self._ssm = create_ssm_executor( + self._ssm_executor = create_ssm_executor( config, request_capacity, mtp_step, @@ -78,14 +78,14 @@ def __init__( @property def ssm_state(self): - return self._ssm.state + return self._ssm_executor.state @property def ssm_slots_per_req(self): - return self._ssm.slots_per_request + return self._ssm_executor.slots_per_request def _reset_tracking(self, req_idx): - self._ssm.clear_history(req_idx) + self._ssm_executor.clear_history(req_idx) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].zero_() @@ -96,7 +96,7 @@ def reset(self, req_idx): self.ssm_state[:, start : start + self.ssm_slots_per_req].zero_() def prepare_prefill(self, reqs): - self._ssm.merge_accepted_updates(reqs) + self._ssm_executor.merge_accepted_updates(reqs) return reqs, reqs * self.ssm_slots_per_req def prepare_decode(self, reqs, cu_seqlens=None): @@ -107,10 +107,12 @@ def prepare_decode(self, reqs, cu_seqlens=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) + self._ssm_executor.prepare_decode(reqs, cu_seqlens) if not speculative: - self._ssm.accept_single_token(reqs) - return SSMForward(self, self._ssm, reqs, self._ssm.decode_indices(reqs, speculative), cu_seqlens, accepted) + self._ssm_executor.accept_single_token(reqs) + return SSMForward( + self, self._ssm_executor, reqs, self._ssm_executor.decode_indices(reqs, speculative), cu_seqlens, accepted + ) def get_accepted_state(self, req_idx): """Return canonical tensors; native MTP gathers and replay materialization allocate snapshots.""" @@ -120,12 +122,12 @@ def get_accepted_state(self, req_idx): 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)) + return ConvSSMTensors(conv, self._ssm_executor.get_accepted_state(req_idx, accepted)) def export_accepted_state(self, req_indices, targets): """Write disjoint targets on the current stream without full-state scratch or changing source history.""" export_indexed_state(self.conv_state, req_indices, [t.conv for t in targets], self.mtp_state_index) - self._ssm.export_accepted_state(req_indices, [t.ssm for t in targets], self.mtp_state_index) + self._ssm_executor.export_accepted_state(req_indices, [t.ssm for t in targets], self.mtp_state_index) def prepare_restore(self, req_idx): @@ -220,7 +222,7 @@ def save_bounded_att_runtime_to_pages( replay_flush = flush_history and self.mtp_state_index is None targets = [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] export_indexed_state(self.conv_state, req_indices, [target.conv for target in targets], self.mtp_state_index) - self._ssm.export_accepted_state( + self._ssm_executor.export_accepted_state( req_indices, [target.ssm for target in targets], self.mtp_state_index, @@ -253,14 +255,16 @@ def validate_bounded_att_commit(self, acceptance: StateAcceptance): acceptance.req_indices.device == self.ssm_state.device ), "commit inputs must match the state device; CUDA states require CUDA tensors" if self.mtp_state_index is not None: - validator = getattr(self._ssm, "validate_commit", None) + validator = getattr(self._ssm_executor, "validate_commit", None) return validator(acceptance, self.state_name) if validator is not None else None return None def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: commit_data = self.validate_bounded_att_commit(acceptance) if self.mtp_state_index is not None: - self._ssm.commit_validated(acceptance, self.mtp_state_index, self.state_name, commit_data=commit_data) + self._ssm_executor.commit_validated( + acceptance, self.mtp_state_index, self.state_name, commit_data=commit_data + ) # ======================================================================== # 三、CPU Cache 数据交换 diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py index 3070856c7a..a0883c305e 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/base.py +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -18,7 +18,7 @@ class SSMExecutor: """ # State 负责资源生命周期和层映射，Executor 只负责本模式的 decode、接受和快照导出。 - # 这样 Native、Compact、Replay 共享同一条上层调用链，差异收敛在执行器内部。 + # 这样 Native、Replay 共享同一条上层调用链，差异收敛在执行器内部。 slots_per_request = 1 _run_config: SSMRunConfig | None = None @@ -39,7 +39,7 @@ def clear_history(self, req): return None def prepare_decode(self, reqs, cu_seqlens=None): - """在 decode 前准备本轮基准；Native/Compact 默认不需要准备。""" + """在 decode 前准备本轮基准；Native 默认不需要准备。""" return None 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 86c5aa9282..7fc7981d32 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -79,9 +79,9 @@ def test_hybrid_layer_mapping_and_mtp_states( 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 + assert bounded_state._ssm_executor.state is bounded_state.ssm_state if mode == "replay": - assert bounded_state._ssm.capacity == 8 + assert bounded_state._ssm_executor.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/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 5b0f4b249f..cd4927361e 100644 --- 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 @@ -29,7 +29,7 @@ def make(n, width): ssm_backend.verify_width = width manager = SimpleNamespace() manager.mtp_state_index = tensor([-9] * (hold + 1)) - manager._ssm = ssm_backend + manager._ssm_executor = ssm_backend return manager, args, (ids, lengths, starts) @@ -66,7 +66,7 @@ def record(*args, **kwargs): def accept_unfused(m, args): update(m.mtp_state_index, *args) - m._ssm.accept_updates(args[1][args[0].long()], m.mtp_state_index) + m._ssm_executor.accept_updates(args[1][args[0].long()], m.mtp_state_index) def accept_fused(m, args): @@ -74,7 +74,7 @@ def accept_fused(m, args): ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) lengths = cumulative[ends.long()] - cumulative[args[0].long()] - m._ssm.commit(StateAcceptance(reqs, lengths), m.mtp_state_index, "main") + m._ssm_executor.commit(StateAcceptance(reqs, lengths), m.mtp_state_index, "main") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -95,15 +95,15 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): 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) + m._ssm_executor.cursors.fill_(7) + ref._ssm_executor.cursors.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): count = (step + i) % (l + 1) mask.extend([1] * count + [0] * (l - count)) expected_index[r] = count - 1 - if r != m._ssm.hold_req_id: + if r != m._ssm_executor.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) accept_unfused(ref, args) @@ -112,12 +112,12 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): 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 torch.equal(m._ssm_executor.cursors, ref._ssm_executor.cursors) assert m.mtp_state_index.tolist() == expected_index - assert m._ssm.cursors.tolist() == expected_cursor + assert m._ssm_executor.cursors.tolist() == expected_cursor checks += 1 - m._ssm = NativeSSMExecutor(torch.empty(0), width) + m._ssm_executor = NativeSSMExecutor(torch.empty(0), width) accept_fused(m, args) assert m.mtp_state_index.tolist() == expected_index @@ -145,7 +145,7 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput managers = [] for ssm_backend in ssm_backends: manager = SimpleNamespace() - manager._ssm = ssm_backend + manager._ssm_executor = ssm_backend manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") managers.append(manager) reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") @@ -214,8 +214,8 @@ def cycle(i): 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._ssm_executor.clear_history(1) + manager._ssm_executor.state[:, 1].zero_() manager.mtp_state_index[1] = 0 for ssm_backend in ssm_backends: ssm_backend.merge_accepted_updates(reqs[:3]) 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 index e88c021fbf..568546d602 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -497,7 +497,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg state.mtp_state_index[1] = 1 state.conv_state.normal_() state.ssm_state.zero_() - backend = state._ssm + backend = state._ssm_executor backend.raw_keys.normal_() backend.raw_values.normal_() backend.gates.uniform_(-0.1, -0.01) 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 index c3efd52b1f..42a0899d7f 100644 --- 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 @@ -102,7 +102,7 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle for cls in [ReplaySSMExecutor, ReplaySSMExecutor] ] manager = SimpleNamespace() - manager._ssm = ssm_backends[1] + manager._ssm_executor = ssm_backends[1] manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) mtp = torch.zeros_like(rows) @@ -150,7 +150,7 @@ def forward_accept(index, cu, reqs): if index == 0: ssm_backend.accept_updates(reqs, accepted) else: - manager._ssm.commit( + manager._ssm_executor.commit( StateAcceptance(rows[starts.long()], accepted[rows[starts.long()]] + 1), manager.mtp_state_index, "main", diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index f4399044e3..cf8cb16c71 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -20,8 +20,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): 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 + if state._ssm_executor is not None: + assert state._ssm_executor.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) @@ -32,8 +32,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): 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 + assert state._ssm_executor.cursors[2] == 0 + assert state._ssm_executor.cursors[1] == 1 torch.cuda.synchronize() @@ -81,7 +81,7 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): 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._ssm_executor is state._ssm_executor assert copy.mtp_state_index is state.mtp_state_index assert all(vars(state)[name] is value for name, value in before.items()) @@ -100,11 +100,11 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( 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 + state._ssm_executor.raw_keys.zero_() + state._ssm_executor.raw_values.zero_() + state._ssm_executor.gates.fill_(-0.5) + state._ssm_executor.betas.zero_() + state._ssm_executor.cursors[1], state._ssm_executor.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() diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index e1adea50e9..a5c47924dc 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -45,11 +45,11 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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 + state._ssm_executor.raw_keys.zero_() + state._ssm_executor.raw_values.zero_() + state._ssm_executor.gates.fill_(-0.5) + state._ssm_executor.betas.zero_() + state._ssm_executor.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") @@ -62,7 +62,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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 + assert state._ssm_executor.cursors[src].item() == 1 for restore in ( lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), @@ -72,7 +72,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if state.mtp_state_index is not None: state.mtp_state_index[dst] = step if mode == "replay": - state._ssm.cursors[dst] = 3 + state._ssm_executor.cursors[dst] = 3 restore() actual = state.get_accepted_state(dst) torch.testing.assert_close(actual.conv, expected_conv, rtol=0, atol=0) @@ -80,7 +80,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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 + assert state._ssm_executor.cursors[dst].item() == 0 state.reset_bounded_att_state([dst]) assert not state.conv_state[:, dst].any() start = dst * state.ssm_slots_per_req @@ -105,21 +105,21 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty state.ssm_state[:, slot].fill_(slot + 1) state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) 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.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + state._ssm_executor.raw_keys.zero_() + state._ssm_executor.raw_values.zero_() + state._ssm_executor.gates.fill_(-0.5) + state._ssm_executor.betas.zero_() + state._ssm_executor.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) requests = [2, 0, 1] expected = [(s.conv.clone(), s.ssm.clone()) for s in map(state.get_accepted_state, requests)] before = state.ssm_state.clone() - cursors = state._ssm.cursors.clone() if mode == "replay" else None + cursors = state._ssm_executor.cursors.clone() if mode == "replay" else None def no_snapshot(*args): raise AssertionError("batch export must write directly to the destination") monkeypatch.setattr(state, "get_accepted_state", no_snapshot) - monkeypatch.setattr(state._ssm, "get_accepted_state", no_snapshot) + monkeypatch.setattr(state._ssm_executor, "get_accepted_state", no_snapshot) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): @@ -131,7 +131,7 @@ def no_snapshot(*args): torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) assert torch.equal(state.ssm_state, before) if cursors is not None: - assert torch.equal(state._ssm.cursors, cursors) + assert torch.equal(state._ssm_executor.cursors, cursors) target, _, _ = make_state(cls, "native", 2, dtype) target.mtp_state_index.fill_(2) target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) @@ -141,7 +141,7 @@ def no_snapshot(*args): torch.testing.assert_close(actual.ssm, ssm, atol=0, rtol=0) assert torch.equal(state.ssm_state, before) if cursors is not None: - assert torch.equal(state._ssm.cursors, cursors) + assert torch.equal(state._ssm_executor.cursors, cursors) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index b519b78519..ca4dc9d3c1 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -86,7 +86,7 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): def test_kda_backend_receives_gate_bound_at_construction(linear_config): owner = KDAState(linear_config, 3, mtp_step=2, ssm_state_mode="replay", gate_lower_bound=-2.5, device="cpu") - assert owner._ssm.gate_lower_bound == -2.5 + assert owner._ssm_executor.gate_lower_bound == -2.5 def test_replay_backend_publishes_detached_run_config(): @@ -110,7 +110,7 @@ def test_replay_backend_publishes_detached_run_config(): @pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): owner = cls(linear_config, request_capacity=2, mtp_step=mtp_step, ssm_state_mode="replay", device="cpu") - assert owner._ssm.compute_mode == expected + assert owner._ssm_executor.compute_mode == expected def make_acceptance(batches=()): @@ -161,10 +161,10 @@ def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") calls = [] monkeypatch.setattr(base, "commit_ssm_prefix", lambda *args: calls.append(args)) - monkeypatch.setattr(owner._ssm, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + monkeypatch.setattr(owner._ssm_executor, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) owner.commit_bounded_att_state(make_acceptance()) assert len(calls) == 1 - assert calls[0][3] is owner._ssm.cursors + assert calls[0][3] is owner._ssm_executor.cursors assert calls[0][0] is owner.mtp_state_index @@ -178,7 +178,7 @@ def test_graph_layout_mismatch_rejected_before_copy(change): if change == "owner": source.resource = copy(owner) elif change == "executor": - source.executor = copy(owner._ssm) + source.executor = copy(owner._ssm_executor) elif change == "shape": source.state_indices = source.state_indices[:1] elif change == "optional": @@ -245,9 +245,9 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): pytest.skip("CUDA required") owner, _, _ = make_state(mode="replay", step=3, device=device) owner.mtp_state_index.fill_(2) - owner._ssm.cursors.fill_(8) + owner._ssm_executor.cursors.fill_(8) reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] owner.commit_bounded_att_state(StateAcceptance(reqs, lengths)) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] - assert owner._ssm.cursors.tolist() == [11, 8, 9, 8] + assert owner._ssm_executor.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 0979efa560..f11647a5e9 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -147,13 +147,13 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( 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 + linear._ssm_executor.raw_keys.zero_() + linear._ssm_executor.raw_values.zero_() + linear._ssm_executor.gates.fill_(-0.5) + linear._ssm_executor.betas.zero_() + linear._ssm_executor.cursors[0] = 1 + ssm = linear._ssm_executor.get_accepted_state(0).clone() + linear._ssm_executor.cursors[2] = 3 if mtp_step: linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( @@ -177,7 +177,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) assert torch.equal(linear.ssm_state[:, 2 * linear.ssm_slots_per_req], ssm) if mode == "replay": - assert linear._ssm.cursors[2] == 0 + assert linear._ssm_executor.cursors[2] == 0 if mtp_step: assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 4eafaad5d4..fdcb28beab 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -256,10 +256,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod if mtp_step: accepted = torch.zeros(4, dtype=torch.int32, device="cuda") accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 - selected._ssm.accept_updates(ids, accepted) + selected._ssm_executor.accept_updates(ids, accepted) for req, length in zip(requests, lengths): torch.testing.assert_close( - selected._ssm.get_accepted_state(req)[0], + selected._ssm_executor.get_accepted_state(req)[0], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index ed219d2ede..7a3884a58d 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -167,7 +167,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) mem.linear.ssm_state[:, src_req * mem.linear.ssm_slots_per_req].copy_(ssm[:, heads]) if source_mode == "replay": - backend = mem.linear._ssm + backend = mem.linear._ssm_executor backend.raw_keys.zero_() backend.raw_values.zero_() backend.gates.fill_(-0.5) @@ -188,7 +188,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() if dest_mode == "replay": - mem.linear._ssm.cursors[dst_req] = 3 + mem.linear._ssm_executor.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: @@ -260,7 +260,7 @@ def test_pd_kv_and_runtime_state_roundtrip( 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 + assert mem.linear._ssm_executor.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( From fec53e8d7c2ce6d546c6a8b41bbf9bf5fc309211 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 8 Nov 2026 20:03:33 +0800 Subject: [PATCH 04/54] style: modernize ReplaySSM Python interfaces --- .../bounded_att/conv_ssm_state.py | 32 +++++++++---------- .../attention_state/bounded_att/interface.py | 32 +++++++++---------- .../attention_state/bounded_att/ssm/export.py | 13 ++++++-- .../bounded_att/ssm/forward.py | 30 +++++++++++------ lightllm/common/attention_state/commit.py | 12 +++---- lightllm/common/attention_state/manager.py | 16 ++++------ lightllm/server/core/objs/start_args_type.py | 2 +- 7 files changed, 78 insertions(+), 59 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py index 3726748512..6f9ae50f97 100644 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -1,7 +1,7 @@ from __future__ import annotations from functools import lru_cache -from typing import TYPE_CHECKING, List, NamedTuple +from typing import TYPE_CHECKING, NamedTuple from math import prod import torch @@ -195,7 +195,7 @@ def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Ten 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: + 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: @@ -203,10 +203,10 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], flush_history: bool = False, ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" @@ -231,10 +231,10 @@ def save_bounded_att_runtime_to_pages( def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], ) -> None: """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance( @@ -283,9 +283,9 @@ def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: def write_state_pages_to_cpu( self, page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], + cpu_page_indices: list[int], layout: StateTransferLayout, cpu_page_ready_flags: list[bool] | None = None, ) -> None: @@ -318,8 +318,8 @@ def write_state_pages_to_cpu( def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], + cpu_page_indices: list[int], + big_page_slot_ids: list[int], layout: StateTransferLayout, ) -> None: """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" @@ -349,8 +349,8 @@ def read_state_pages_from_cpu( def copy_big_pages_from_dp( self, - src_page_ids: List[int], - dst_page_ids: List[int], + src_page_ids: list[int], + dst_page_ids: list[int], remote_dp_bounded_state: _ConvSSMState, ) -> None: """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" @@ -373,8 +373,8 @@ def copy_big_pages_from_dp( def copy_runtime_state_from_dp( self, - src_req_indices: List[int], - dst_req_indices: List[int], + src_req_indices: list[int], + dst_req_indices: list[int], remote_dp_bounded_state: _ConvSSMState, ) -> None: """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index 983d8e78c5..468fd0da28 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from functools import lru_cache -from typing import Any, List +from typing import Any import torch @@ -65,17 +65,17 @@ def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod - def reset_bounded_att_state(self, req_indices: List[int]) -> None: + def reset_bounded_att_state(self, req_indices: list[int]) -> None: """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" ... @abstractmethod def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], flush_history: bool = False, ) -> None: """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 @@ -89,10 +89,10 @@ def save_bounded_att_runtime_to_pages( @abstractmethod def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], ) -> None: """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 @@ -129,9 +129,9 @@ def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: def write_state_pages_to_cpu( self, page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], + cpu_page_indices: list[int], layout: StateTransferLayout, cpu_page_ready_flags: list[bool] | None = None, ) -> None: @@ -142,8 +142,8 @@ def write_state_pages_to_cpu( def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], + cpu_page_indices: list[int], + big_page_slot_ids: list[int], layout: StateTransferLayout, ) -> None: """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" @@ -158,8 +158,8 @@ class DpCopyInterface(ABC): @abstractmethod def copy_big_pages_from_dp( self, - src_page_ids: List[int], - dst_page_ids: List[int], + src_page_ids: list[int], + dst_page_ids: list[int], remote_dp_bounded_state: CoreInterface, ) -> None: """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" @@ -168,8 +168,8 @@ def copy_big_pages_from_dp( @abstractmethod def copy_runtime_state_from_dp( self, - src_req_indices: List[int], - dst_req_indices: List[int], + src_req_indices: list[int], + dst_req_indices: list[int], remote_dp_bounded_state: CoreInterface, ) -> None: diff --git a/lightllm/common/attention_state/bounded_att/ssm/export.py b/lightllm/common/attention_state/bounded_att/ssm/export.py index f04c4313c5..bfb55bc535 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/export.py +++ b/lightllm/common/attention_state/bounded_att/ssm/export.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from math import prod import torch @@ -7,7 +8,7 @@ from ..copy_kernel import copy_bounded_att_tensors -def upload_indices(values, device, dtype=torch.int64): +def upload_indices(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.int64) -> torch.Tensor: """把 Python 索引上传到目标 device，使用 pinned memory 减少同步。""" return torch.tensor(values, dtype=dtype, pin_memory=True).to(device, non_blocking=True) @@ -49,7 +50,15 @@ def _export_indexed( tl.store(target + target_offset, values, valid) -def export_indexed_state(source, req_indices, targets, accepted_indices=None, *, slots_per_request=1, index_stride=1): +def export_indexed_state( + source: torch.Tensor, + req_indices: Sequence[int], + targets: Sequence[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + slots_per_request: int = 1, + index_stride: int = 1, +) -> None: """将 source 中的 accepted state 写入预分配的 target 页。 ``req_indices`` 和 ``targets`` 必须一一对应；MTP 通过 diff --git a/lightllm/common/attention_state/bounded_att/ssm/forward.py b/lightllm/common/attention_state/bounded_att/ssm/forward.py index 6930630d53..a23054299a 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/forward.py +++ b/lightllm/common/attention_state/bounded_att/ssm/forward.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from dataclasses import dataclass, field -from typing import Optional, TYPE_CHECKING +from typing import TYPE_CHECKING import torch @@ -21,16 +23,26 @@ class SSMForward: 后续层复用同一配置，``finish`` 再把它封装进 ``StateCommit``。 """ - resource: "_ConvSSMState" - executor: "SSMExecutor" + resource: _ConvSSMState + executor: SSMExecutor req_indices: torch.Tensor state_indices: torch.Tensor - cu_seqlens: Optional[torch.Tensor] - accepted_token_counts: Optional[torch.Tensor] - run_config: Optional[SSMRunConfig] = None + cu_seqlens: torch.Tensor | None + accepted_token_counts: torch.Tensor | None + run_config: SSMRunConfig | None = None _layout_checked: bool = field(default=False, init=False, repr=False) - def run_layer(self, logical_layer, q, k, v, a, b, a_log, bias): + def run_layer( + self, + logical_layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + ) -> torch.Tensor: """执行一层线性注意力；logical_layer 会映射到物理权重层。""" output, self.run_config = self.executor.forward( @@ -50,7 +62,7 @@ def run_layer(self, logical_layer, q, k, v, a, b, a_log, bias): ) return output - def finish(self): + def finish(self) -> StateCommit | None: """结束 forward，返回接受阶段需要的状态提交描述。""" payload = self.executor.create_commit(self.run_config) @@ -58,7 +70,7 @@ def finish(self): return None return StateCommit(self.resource.state_name, payload) - def copy_inputs(self, source: "SSMForward"): + def copy_inputs(self, source: SSMForward) -> None: """把动态图输入复制到已 capture 的 tensor。 graph bucket 的资源和 tensor layout 在首次绑定时校验一次，后续调用 diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py index 6256ea2c80..f9dc147975 100644 --- a/lightllm/common/attention_state/commit.py +++ b/lightllm/common/attention_state/commit.py @@ -1,5 +1,5 @@ +from collections.abc import Iterator from dataclasses import dataclass -from typing import Tuple import torch @@ -17,9 +17,9 @@ class StateCommitBatch: """描述一个 microbatch 分区及其各状态的 forward 配置。""" request_count: int - commits: Tuple[StateCommit, ...] = () + commits: tuple[StateCommit, ...] = () - def __post_init__(self): + def __post_init__(self) -> None: 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] @@ -41,9 +41,9 @@ class StateAcceptance: req_indices: torch.Tensor accepted_lengths: torch.Tensor - batches: Tuple[StateCommitBatch, ...] = () + batches: tuple[StateCommitBatch, ...] = () - def validate(self, state_names=None): + def validate(self, state_names: set[str] | None = None) -> None: """检查请求、接受长度和 forward 分区是否能组成一次提交。""" if self.req_indices.ndim != 1 or self.accepted_lengths.shape != self.req_indices.shape: @@ -58,7 +58,7 @@ def validate(self, state_names=None): 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): + def partitions(self, state_name: str) -> Iterator[tuple[torch.Tensor, object | None]]: """按 forward 分区返回请求索引和指定状态的提交配置。""" start = 0 diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 5bb052374a..bccb3c63e6 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -2,8 +2,6 @@ from typing import TYPE_CHECKING from functools import lru_cache -from typing import List - import torch from .commit import StateAcceptance @@ -333,7 +331,7 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # ======================================================================== # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 - def reset_bounded_att_state(self, req_indices: List[int]) -> None: + def reset_bounded_att_state(self, req_indices: list[int]) -> None: """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" for state in self.get_bounded_attention_states(): state.reset_bounded_att_state(req_indices) @@ -342,10 +340,10 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], flush_history: bool = False, ) -> None: """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" @@ -370,10 +368,10 @@ def save_bounded_att_runtime_to_pages( def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], ) -> None: """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" bounded_states = self.get_bounded_attention_states() diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9cf0f719ab..0609132e62 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -256,7 +256,7 @@ class StartArgs: # hybrid attention model (Qwen3Next) ssm_state_mode: str = "native" - ssm_replay_config: dict = field(default_factory=dict) + ssm_replay_config: dict[str, object] = field(default_factory=dict) linear_att_hash_page_size: int = field(default=512) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) From f88e7ba0de29f38c15510c9a60175cf136a42397 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 8 Nov 2026 20:50:22 +0800 Subject: [PATCH 05/54] style: document ReplaySSM Python functions --- .../bounded_att/conv_ssm_state.py | 59 +++++---- .../bounded_att/gdn/create_gdn_state.py | 2 +- .../bounded_att/kda/create_kda_state.py | 2 +- .../bounded_att/ssm/__init__.py | 21 ++- .../attention_state/bounded_att/ssm/base.py | 83 +++++++----- .../attention_state/bounded_att/ssm/native.py | 49 ++++--- .../attention_state/bounded_att/ssm/replay.py | 121 +++++++++++------- .../basemodel/attention/linear/state.py | 32 +++-- .../linear_att/ssm_acceptance.py | 11 +- .../triton_kernel/linear_att/ssm_autotune.py | 44 ++++--- lightllm/common/ssm_types.py | 3 +- lightllm/models/glm5_next/config.py | 3 +- 12 files changed, 283 insertions(+), 147 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py index 6f9ae50f97..62e3e83399 100644 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -1,7 +1,7 @@ from __future__ import annotations from functools import lru_cache -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple from math import prod import torch @@ -46,9 +46,10 @@ def __init__( device: str | torch.device = "cuda", *, ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, + ssm_replay_config: dict[str, object] | None = None, gate_lower_bound: float = -5.0, - ): + ) -> None: + """创建共享 Conv/SSM 运行态，并按模式绑定对应的 SSM executor。""" self.config = config if self.state_mode is not None and ssm_state_mode != self.state_mode: raise ValueError(f"state mode {ssm_state_mode!r} does not match {self.state_mode!r}") @@ -77,29 +78,35 @@ def __init__( self.small_pages_pool = None @property - def ssm_state(self): + def ssm_state(self) -> torch.Tensor: + """返回 executor 持有的 canonical SSM state。""" return self._ssm_executor.state @property - def ssm_slots_per_req(self): + def ssm_slots_per_req(self) -> int: + """返回每个请求占用的物理 SSM 槽位数。""" return self._ssm_executor.slots_per_request - def _reset_tracking(self, req_idx): + def _reset_tracking(self, req_idx: int) -> None: + """清理请求的临时 history 和 MTP 接受游标。""" self._ssm_executor.clear_history(req_idx) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].zero_() - def reset(self, req_idx): + def reset(self, req_idx: int) -> None: + """清空请求的 Conv/SSM 运行态，同时保留页池快照。""" 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): + def prepare_prefill(self, reqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """合并已接受 history，并返回 prefill 使用的请求和物理索引。""" self._ssm_executor.merge_accepted_updates(reqs) return reqs, reqs * self.ssm_slots_per_req - def prepare_decode(self, reqs, cu_seqlens=None): + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> SSMForward: + """准备普通或 speculative decode，并返回对应的 SSMForward 上下文。""" speculative = cu_seqlens is not None accepted = None if speculative: @@ -114,8 +121,8 @@ def prepare_decode(self, reqs, cu_seqlens=None): self, self._ssm_executor, reqs, self._ssm_executor.decode_indices(reqs, speculative), cu_seqlens, accepted ) - def get_accepted_state(self, req_idx): - """Return canonical tensors; native MTP gathers and replay materialization allocate snapshots.""" + def get_accepted_state(self, req_idx: int) -> ConvSSMTensors: + """返回 canonical Conv/SSM 状态；Native 会选槽位，Replay 会物化快照。""" 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] @@ -124,12 +131,13 @@ def get_accepted_state(self, req_idx): conv = self.conv_state[:, req_idx].index_select(-1, offsets) return ConvSSMTensors(conv, self._ssm_executor.get_accepted_state(req_idx, accepted)) - def export_accepted_state(self, req_indices, targets): - """Write disjoint targets on the current stream without full-state scratch or changing source history.""" + def export_accepted_state(self, req_indices: list[int], targets: list[ConvSSMTensors]) -> None: + """在当前 stream 写入互不重叠的目标页，不分配整批临时快照。""" export_indexed_state(self.conv_state, req_indices, [t.conv for t in targets], self.mtp_state_index) self._ssm_executor.export_accepted_state(req_indices, [t.ssm for t in targets], self.mtp_state_index) - def prepare_restore(self, req_idx): + def prepare_restore(self, req_idx: int) -> ConvSSMTensors: + """清理请求 tracking，并返回可写入恢复快照的 Conv/SSM 视图。""" self._reset_tracking(req_idx) return ConvSSMTensors( @@ -137,15 +145,18 @@ def prepare_restore(self, req_idx): self.ssm_state[:, req_idx * self.ssm_slots_per_req], ) - def restore(self, req_idx, conv, ssm): + def restore(self, req_idx: int, conv: torch.Tensor, ssm: torch.Tensor) -> None: + """把 Conv/SSM 快照复制回指定请求并完成恢复后的 tracking 初始化。""" 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): + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """把逻辑层号映射到当前 State 的物理层号。""" return self._logical_to_storage_layer[logical_layer_id] def get_state_type(self) -> str: + """返回包含算法和 Conv/SSM 精度的稳定 State 类型名。""" return ( f"{self.algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" @@ -215,10 +226,9 @@ def save_bounded_att_runtime_to_pages( ), "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) - # Hybrid radix pages are durable checkpoints. ReplaySSM can fold the - # accepted ring into the live canonical state and the page in one - # kernel, so the next decode starts from an empty history. MTP keeps - # a separate Conv candidate layout and retains the old behavior. + # Hybrid radix page 是持久 checkpoint。ReplaySSM 可以在一个 kernel 中把 + # accepted ring 折叠到 canonical state 和目标页，使下一轮 decode 从空 history 开始。 + # MTP 使用独立的 Conv candidate 布局，因此保留原有行为。 replay_flush = flush_history and self.mtp_state_index is None targets = [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] export_indexed_state(self.conv_state, req_indices, [target.conv for target in targets], self.mtp_state_index) @@ -250,7 +260,8 @@ def restore_bounded_att_runtime_from_pages( copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def validate_bounded_att_commit(self, acceptance: StateAcceptance): + def validate_bounded_att_commit(self, acceptance: StateAcceptance) -> object | None: + """校验 acceptance 设备并委托 executor 检查可选的提交配置。""" assert ( acceptance.req_indices.device == self.ssm_state.device ), "commit inputs must match the state device; CUDA states require CUDA tensors" @@ -260,6 +271,7 @@ def validate_bounded_att_commit(self, acceptance: StateAcceptance): return None def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: + """应用已接受的 MTP 前缀并推进对应的 executor 状态。""" commit_data = self.validate_bounded_att_commit(acceptance) if self.mtp_state_index is not None: self._ssm_executor.commit_validated( @@ -562,7 +574,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: ) return tuple(regions) - def _pool(self, page_kind: PageKind): + def _pool(self, page_kind: PageKind) -> Any: """根据大小页类型取得已启用的页池，不访问 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 @@ -570,7 +582,8 @@ def _pool(self, page_kind: PageKind): return pool @staticmethod - def _fields(conv: torch.Tensor, ssm: torch.Tensor): + def _fields(conv: torch.Tensor, ssm: torch.Tensor) -> dict[str, TensorField]: + """把 Conv/SSM 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 35eb5147e3..0ba7d1f112 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 @@ -15,7 +15,7 @@ def create_gdn_state( device: str | torch.device = "cuda", *, ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, + ssm_replay_config: dict[str, object] | None = None, gate_lower_bound: float = -5.0, ) -> GDNState: """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" 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 ba4614da27..af90adad00 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 @@ -15,7 +15,7 @@ def create_kda_state( device: str | torch.device = "cuda", *, ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, + ssm_replay_config: dict[str, object] | None = None, gate_lower_bound: float = -5.0, ) -> KDAState: """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" diff --git a/lightllm/common/attention_state/bounded_att/ssm/__init__.py b/lightllm/common/attention_state/bounded_att/ssm/__init__.py index 15dcd1501a..3cf73492ec 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/__init__.py +++ b/lightllm/common/attention_state/bounded_att/ssm/__init__.py @@ -1,7 +1,26 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + from ..conv_ssm_types import ConvSSMStateConfig + from .base import SSMExecutor + -def create_ssm_executor(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound): +def create_ssm_executor( + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int, + device: str | torch.device, + mode: str, + replay_config: dict[str, object] | None, + *, + kda: bool, + gate_lower_bound: float, +) -> SSMExecutor: + """根据状态模式创建 Native 或 Replay SSM executor。""" """按统一 state layout 创建 native 或 replay executor。 native 为每个 MTP 位置分配物理槽位，replay 为每个请求分配一个槽位， diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py index a0883c305e..68c2adcc7a 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/base.py +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -1,20 +1,28 @@ +from __future__ import annotations + +import torch + +from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.ssm_types import SSMRunConfig from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from .export import export_indexed_state -def update_accepted_indices(acceptance, indices, *, replay_cursors=None): +def update_accepted_indices( + acceptance: StateAcceptance, indices: torch.Tensor, *, replay_cursors: torch.Tensor | None = None +) -> None: + """把接受长度写入 MTP 槽位索引，并按需推进 Replay cursor。""" + commit_ssm_prefix(indices, acceptance.req_indices, acceptance.accepted_lengths, replay_cursors) class SSMExecutor: """一种 SSM 模式的 decode、接受和快照适配器。 - the algorithm-specific state 持有资源和逻辑层到物理层的映射；executor 只处理 - decode 递推、接受前缀以及 accepted-state 导出。prefill 仍由 attention - backend 执行。一次 verify 的布局由 ``SSMForward`` 保存，并在 commit - 阶段通过 ``StateCommit`` 传回同一个 executor。 + 算法 State 持有资源和逻辑层到物理层的映射；executor 只处理 decode 递推、接受前缀 + 以及 accepted-state 导出。prefill 仍由 attention backend 执行。一次 verify 的 + 布局由 ``SSMForward`` 保存，并在 commit 阶段通过 ``StateCommit`` 传回同一个 executor。 """ # State 负责资源生命周期和层映射，Executor 只负责本模式的 decode、接受和快照导出。 @@ -30,78 +38,93 @@ def run_config(self) -> SSMRunConfig | None: return self._run_config @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None): + def run_config(self, config: SSMRunConfig | dict[str, int] | None) -> None: + """保存不可变的 kernel 配置，兼容字典形式的 autotune 结果。""" self._run_config = config if config is None or isinstance(config, SSMRunConfig) else SSMRunConfig(**config) - def clear_history(self, req): + def clear_history(self, req: int | torch.Tensor) -> None: """清空请求的临时 history；默认模式没有额外 history。""" return None - def prepare_decode(self, reqs, cu_seqlens=None): + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> None: """在 decode 前准备本轮基准；Native 默认不需要准备。""" return None - def accept_single_token(self, reqs): + def accept_single_token(self, reqs: torch.Tensor) -> None: """提交非 speculative decode 产生的单 token 更新。""" return None - def merge_accepted_updates(self, reqs): + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把已接受的临时更新合并回 canonical state。""" return None - def decode_indices(self, reqs, speculative): + def decode_indices(self, reqs: torch.Tensor, speculative: bool) -> torch.Tensor: """返回 decode 使用的物理槽位；Native MTP 会展开请求槽位。""" return reqs - def get_accepted_state(self, req_idx, accepted_index=None): + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: """返回一个请求的 accepted-state 视图或物化快照。""" return self.state[:, req_idx] - def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: """把 accepted state 写入调用方提供的目标页，不负责分配目标。""" export_indexed_state(self.state, req_indices, targets) - def create_commit(self, run_config): + def create_commit(self, run_config: SSMRunConfig | None) -> object | None: """把本次 forward 的配置转换为接受阶段可携带的提交 payload。""" return None - def commit(self, acceptance, indices, state_name): + def commit(self, acceptance: StateAcceptance, indices: torch.Tensor, state_name: str) -> None: """校验并应用接受结果；供直接调用 State 接口时使用。""" validator = getattr(self, "validate_commit", None) commit_data = validator(acceptance, state_name) if validator is not None else None self.commit_validated(acceptance, indices, state_name, commit_data=commit_data) - def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): + def commit_validated( + self, + acceptance: StateAcceptance, + indices: torch.Tensor, + state_name: str, + *, + commit_data: object | None = None, + ) -> None: """应用已由 State/Manager 校验过的接受结果，不重复检查公共契约。""" update_accepted_indices(acceptance, indices) def forward( self, - layer, - q, - k, - v, - a, - b, - a_log, - bias, - reqs, - cu_seqlens=None, - run_config=None, + layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, *, - state_indices=None, - accepted_token_counts=None, - ): + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig | None]: """执行一层 decode，并返回输出及本次实际采用的 run config。""" raise NotImplementedError diff --git a/lightllm/common/attention_state/bounded_att/ssm/native.py b/lightllm/common/attention_state/bounded_att/ssm/native.py index 92ee179b04..b664f78027 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/native.py +++ b/lightllm/common/attention_state/bounded_att/ssm/native.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import torch +from lightllm.common.ssm_types import SSMRunConfig from .base import SSMExecutor from .export import export_indexed_state @@ -7,13 +10,16 @@ class NativeSSMExecutor(SSMExecutor): """直接调用 fused recurrent kernel，并把 MTP 状态展开为物理槽位。""" - def __init__(self, state, verify_width, *, kda=False, gate_lower_bound=-5.0): + def __init__( + self, state: torch.Tensor, verify_width: int, *, kda: bool = False, gate_lower_bound: float = -5.0 + ) -> None: + """创建直接写入 canonical state 的 Native executor。""" 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): + def decode_indices(self, reqs: torch.Tensor, speculative: bool) -> torch.Tensor: """返回普通 decode 或 speculative verify 对应的物理 state 索引。""" if not speculative: @@ -21,7 +27,7 @@ def decode_indices(self, reqs, speculative): 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): + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: """返回单槽位 view，或 MTP 接受位置对应的 state 副本。""" if self.slots_per_request == 1: @@ -29,7 +35,14 @@ def get_accepted_state(self, req_idx, accepted_index=None): index = req_idx * self.slots_per_request + accepted_index return self.state.index_select(1, index.long().view(1))[:, 0] - def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: """把 native state 按请求和接受位置导出到调用方目标页。""" export_indexed_state( @@ -43,21 +56,21 @@ def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, def forward( self, - layer, - q, - k, - v, - a, - b, - a_log, - bias, - reqs, - cu_seqlens=None, - run_config=None, + layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, *, - state_indices=None, - accepted_token_counts=None, - ): + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: """执行 native recurrent decode，并原地更新 fused kernel 的最终 state。""" state = self.state[layer] diff --git a/lightllm/common/attention_state/bounded_att/ssm/replay.py b/lightllm/common/attention_state/bounded_att/ssm/replay.py index 3378fb3c04..51987ba61a 100644 --- a/lightllm/common/attention_state/bounded_att/ssm/replay.py +++ b/lightllm/common/attention_state/bounded_att/ssm/replay.py @@ -1,6 +1,9 @@ +from __future__ import annotations + import torch import triton +from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.ssm_types import SSMRunConfig from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _accept_updates, @@ -25,17 +28,17 @@ class ReplaySSMExecutor(SSMExecutor): def __init__( self, - state, - capacity, - verify_width, - activation_dtype=torch.bfloat16, + state: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, *, - num_key_heads=None, - kda=False, - gate_lower_bound=-5.0, - compute_mode="inline", - run_config=None, - ): + num_key_heads: int | None = None, + kda: bool = False, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> None: """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 @@ -58,16 +61,15 @@ def __init__( 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. + # 保留冗余布局，兼容已有调用方和 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). - # BF16 can use a narrower short-history buffer for the precompute path; - # FP32 remains the safe default for all other shapes. + # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 + # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 + # 下可以使用更窄的 history buffer；其他形状使用 FP32 作为安全默认值。 history_dtype = ( activation_dtype if activation_dtype == torch.bfloat16 @@ -88,15 +90,15 @@ def __init__( 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.""" + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" if type(req) is int: - # Scalar assignment copies a CPU tensor into the CUDA slot. + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 self.cursors[req].zero_() else: self.cursors[req] = 0 - def prepare_decode(self, reqs, cu_seqlens=None): + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" positions = torch.empty_like(reqs) layers, _, hv, kd, vd = self.state.shape @@ -146,7 +148,7 @@ def prepare_decode(self, reqs, cu_seqlens=None): ) return positions - def accept_updates(self, reqs, accepted_state_indices=None): + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: """更新 cursor 以保留每个请求本轮已接受的前缀。 MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft @@ -162,14 +164,21 @@ def accept_updates(self, reqs, accepted_state_indices=None): 256, ) - def merge_accepted_updates(self, reqs): + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把接受的 history 物化回 canonical state，并清空请求 cursor。""" 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, output_strides=None, flush_export=False): + def _materialize_accepted_state( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" assert not (snapshot and flush_export), "flush export must update the canonical state" @@ -203,7 +212,7 @@ def _materialize_accepted_state(self, reqs, output, snapshot, output_strides=Non num_stages=config.num_stages, ) - def get_accepted_state(self, req_idx, accepted_index=None): + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: """返回不修改 active state/history 的单请求 canonical 快照。""" layers, _, hv, k, v = self.state.shape output = torch.empty((layers, hv, k, v), dtype=self.state.dtype, device=self.state.device) @@ -211,7 +220,14 @@ def get_accepted_state(self, req_idx, accepted_index=None): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, flush=False): + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" if not req_indices: @@ -228,28 +244,38 @@ def export_accepted_state(self, req_indices, targets, accepted_indices=None, *, output_strides=strides, ) - def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> SSMRunConfig: """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) def forward( self, - linear_layer_idx, - q, - k, - v, - a, - b, - a_log, - bias, - reqs, - cu_seqlens=None, - run_config=None, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, *, - state_indices=None, - accepted_token_counts=None, - ): + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig]: """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) @@ -259,8 +285,8 @@ def forward( out = torch.empty_like(v) bv = config.BV precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode == "precompute" - # Resolve tile dimensions in Python: bare max() in Triton JIT code - # lowers to a tensor operation, which cannot define static shapes. + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor + # 运算，无法定义静态 shape。 bw = triton.next_power_of_2(self.verify_width) _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, @@ -307,16 +333,23 @@ def forward( ) return out, config - def accept_single_token(self, reqs): + def accept_single_token(self, reqs: torch.Tensor) -> None: """非 MTP decode 直接把单 token 写入当前 history。""" self.accept_updates(reqs) - def commit(self, acceptance, indices, state_name): + def commit(self, acceptance: StateAcceptance, indices: torch.Tensor, state_name: str) -> None: """应用接受结果；MTP 直接推进 cursor，单 token 还需写入 history。""" self.commit_validated(acceptance, indices, state_name) - def commit_validated(self, acceptance, indices, state_name, *, commit_data=None): + def commit_validated( + self, + acceptance: StateAcceptance, + indices: torch.Tensor, + state_name: str, + *, + commit_data: object | None = None, + ) -> None: """应用已验证的接受结果，推进 replay cursor。""" if self.verify_width > 1: diff --git a/lightllm/common/basemodel/attention/linear/state.py b/lightllm/common/basemodel/attention/linear/state.py index 1ce0a534b6..a3142e6156 100644 --- a/lightllm/common/basemodel/attention/linear/state.py +++ b/lightllm/common/basemodel/attention/linear/state.py @@ -1,15 +1,21 @@ +from __future__ import annotations + import copy from dataclasses import dataclass, field +from typing import Any import torch +from lightllm.common.attention_state.commit import StateCommit +from lightllm.common.attention_state.bounded_att.ssm.forward import SSMForward 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): +def get_resources(att_state: Any) -> tuple[Any, ...]: + """返回当前角色下与线性 attention 算法匹配的 bounded State。""" return tuple( state for state in att_state.infer_state.att_state_manager.get_bounded_attention_states( @@ -19,7 +25,8 @@ def get_resources(att_state): ) -def get_resource(att_state, logical_layer): +def get_resource(att_state: Any, logical_layer: int) -> Any: + """按逻辑层号取得当前角色对应的 bounded State。""" 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, @@ -32,7 +39,8 @@ def get_resource(att_state, logical_layer): class LinearPrefillAttState(BasePrefillAttState): state_indices: dict = field(default_factory=dict) - def init_state(self): + def init_state(self) -> None: + """为 prefill 建立每个 State 的请求和物理索引。""" self.state_indices = { state.state_name: state.prepare_prefill(self.infer_state.b_req_idx) for state in get_resources(self) @@ -44,10 +52,12 @@ class LinearDecodeAttState(BaseDecodeAttState): forwards: dict = field(default_factory=dict) _layout_checked: bool = field(default=False, init=False, repr=False) - def get_draft_step(self): + def get_draft_step(self) -> int: + """返回当前模型的 MTP draft 步数。""" return self.backend.mtp_step - def init_state(self): + def init_state(self) -> None: + """根据固定或动态 MTP 布局创建每个 State 的 SSMForward。""" self._layout_checked = False draft_step = self.get_draft_step() infer_state = self.infer_state @@ -73,19 +83,23 @@ def init_state(self): 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): + def get_forward(self, logical_layer: int) -> SSMForward: + """返回指定逻辑层共享的 forward 上下文。""" return self.forwards[get_resource(self, logical_layer).state_name] - def get_state_commits(self): + def get_state_commits(self) -> tuple[StateCommit, ...]: + """收集各 State 在本轮 forward 产生的提交凭据。""" return tuple(commit for forward in self.forwards.values() if (commit := forward.finish()) is not None) - def clone_for_warmup(self, infer_state): + def clone_for_warmup(self, infer_state: Any) -> LinearDecodeAttState: + """复制 decode 状态供 CUDA Graph warmup 使用。""" state = super().clone_for_warmup(infer_state) state.forwards = {name: copy.copy(forward) for name, forward in self.forwards.items()} state._layout_checked = False return state - def copy_for_decode_cuda_graph(self, new_state): + def copy_for_decode_cuda_graph(self, new_state: LinearDecodeAttState) -> None: + """把动态图输入复制到已 capture 的 CUDA Graph 状态。""" if not self._layout_checked: if self.forwards.keys() != new_state.forwards.keys(): raise ValueError("captured linear attention resources changed") diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py index 058a0e6d65..d1e1040213 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -1,3 +1,6 @@ +from __future__ import annotations + +import torch import triton import triton.language as tl @@ -24,7 +27,13 @@ def _commit_prefix( tl.store(Cursors + req, cursor + length, valid) -def commit_ssm_prefix(indices, reqs, lengths, replay_cursors=None): +def commit_ssm_prefix( + indices: torch.Tensor, + reqs: torch.Tensor, + lengths: torch.Tensor, + replay_cursors: torch.Tensor | None = None, +) -> None: + """把每个请求的 accepted length 写入 MTP 索引并推进 Replay cursor。""" if not reqs.is_cuda: indices[reqs] = (lengths - 1).to(indices.dtype) if replay_cursors is not None: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index bcfe691d42..c4e78a2b65 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -1,4 +1,6 @@ -"""Tune one layer's SSM update cycle using the shared autotuner.""" +"""使用共享 autotuner 调优单层 SSM update cycle。""" + +from __future__ import annotations import ast @@ -12,7 +14,8 @@ logger = init_logger(__name__) -def get_configs(): +def get_configs() -> list[dict[str, int]]: + """返回 ReplaySSM update cycle 的候选 Triton 配置。""" return [ {"BV": bv, "num_warps": warps, "num_stages": stages} for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] @@ -20,11 +23,11 @@ def get_configs(): ] -def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens): +def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens) -> dict[str, object]: + """构造区分状态布局、数据类型和投影模式的静态调优键。""" 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. + # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 key = { "mode": mode, "v": 6, @@ -46,13 +49,15 @@ def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens): return key -def run_key(ssm_backend, q, cu_seqlens): +def run_key(ssm_backend, q, cu_seqlens) -> tuple[int, int]: + """返回调优使用的 token 数和请求数键。""" 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. +def scratch_size_bytes(ssm_backend, batch, sequences, inputs) -> int: + """估算一次 SSM autotune scratch 所需的字节数。""" + # State/history 的首轴是 layer，cursor 在各 layer 之间共享。 layers, slots = ssm_backend.state.shape[:2] pool_bytes = sum( x.numel() * x.element_size() // (layers if x.ndim > 1 else 1) @@ -67,14 +72,17 @@ def scratch_size_bytes(ssm_backend, batch, sequences, inputs): 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): +def rebuild_inputs( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> tuple[tuple[object, ...], dict[str, object]]: + """创建不影响服务态的独立输入，供 autotune benchmark 使用。""" 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. + # 为 benchmark 中间量、CUDA Graph 和 allocator 碎片保留一半显存。 if required > available // 2: raise torch.cuda.OutOfMemoryError( f"SSM tuning scratch needs {required} bytes, exceeding its {available // 2}-byte budget" @@ -91,7 +99,7 @@ def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=Non gate_lower_bound=ssm_backend.gate_lower_bound, compute_mode=ssm_backend.compute_mode, ) - # Benchmark callbacks must not recursively select a configuration. + # benchmark 回调不能再次触发配置选择。 scratch._config_is_fixed = True inputs = [ torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype).normal_(0, 0.2) @@ -131,15 +139,16 @@ def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=Non 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.""" +def select_config( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> dict[str, int] | SSMRunConfig | None: + """搜索 update cycle 布局；计时只使用重建的临时输入。""" if workload is None: return run_config ssm_backend.run_config = run_config reqs, accepted = workload ssm_backend.cursors.zero_() - # Include an entire history period and the subsequent fold. Repeating - # just the first round would benchmark an empty history forever. + # 覆盖完整 history 周期并触发一次 fold；只重复第一轮会一直测量空 history。 tokens = q.shape[1 if cu_seqlens is not None else 0] rounds = ssm_backend.capacity // min(ssm_backend.verify_width, tokens) + 1 for step in range(rounds): @@ -149,7 +158,8 @@ def select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None return run_config -def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens): +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> SSMRunConfig: + """返回当前 batch 的固定或已调优 SSM kernel 配置。""" default = ssm_backend.run_config or SSMRunConfig(32, 1) if ssm_backend._config_is_fixed: return default @@ -160,7 +170,7 @@ def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens return default config = select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) config = SSMRunConfig(**config) if config else default - # Each captured graph keeps its own layout, even after other buckets warm up. + # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 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/ssm_types.py b/lightllm/common/ssm_types.py index 636f7ceffe..a0d4da1389 100644 --- a/lightllm/common/ssm_types.py +++ b/lightllm/common/ssm_types.py @@ -7,7 +7,8 @@ class SSMRunConfig: num_warps: int num_stages: int = 3 - def __post_init__(self): + def __post_init__(self) -> None: + """校验 Triton update cycle 配置的字段类型。""" 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/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index f19f9028d3..31d79697f6 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -14,7 +14,7 @@ @dataclasses.dataclass class Glm5NextResourceConfig(Qwen3NextResourceConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + """描述复制式 MLA/index KV 与 TP 分片 KDA checkpoint 的资源配置。""" # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 gate_lower_bound: float = -5.0 @@ -24,6 +24,7 @@ class Glm5NextResourceConfig(Qwen3NextResourceConfig): @classmethod def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict) -> "Glm5NextResourceConfig": + """从模型配置和启动参数构造 GLM5.3 Flash 资源配置。""" config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp From 8720368396969eb0fd7205f7e7cbf6d1eba694ec Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 8 Nov 2026 14:18:31 +0000 Subject: [PATCH 06/54] refactor: split GDN and KDA replay into standalone state modules Drop the shared SSM executor hierarchy so each algorithm keeps its own replay implementation and the native paths stay easier to review. Co-authored-by: Cursor --- .../attention_state/bounded_att/__init__.py | 8 +- .../bounded_att/conv_ssm_state.py | 590 --------- .../bounded_att/gdn/__init__.py | 10 +- .../bounded_att/gdn/create_gdn_state.py | 28 +- .../bounded_att/gdn/gdn_state.py | 529 +++++++- .../bounded_att/gdn/replay_gdn_state.py | 1102 +++++++++++++++++ .../attention_state/bounded_att/interface.py | 51 +- .../bounded_att/kda/__init__.py | 10 +- .../bounded_att/kda/create_kda_state.py | 28 +- .../bounded_att/kda/kda_state.py | 529 +++++++- .../bounded_att/kda/replay_kda_state.py | 1102 +++++++++++++++++ .../sliding_window/fp8_e8m0_swa_state.py | 4 +- .../bounded_att/ssm/__init__.py | 54 - .../attention_state/bounded_att/ssm/base.py | 130 -- .../attention_state/bounded_att/ssm/export.py | 104 -- .../bounded_att/ssm/forward.py | 98 -- .../attention_state/bounded_att/ssm/native.py | 130 -- .../attention_state/bounded_att/ssm/replay.py | 359 ------ .../window_compressor_state.py | 4 +- lightllm/common/attention_state/manager.py | 37 +- .../common/basemodel/attention/base_att.py | 11 - .../basemodel/attention/linear/flashqla.py | 4 +- .../common/basemodel/attention/linear/gdn.py | 270 +++- .../common/basemodel/attention/linear/kda.py | 171 ++- .../basemodel/attention/linear/state.py | 108 -- .../basemodel/attention/linear/triton.py | 4 +- lightllm/common/basemodel/basemodel.py | 7 - lightllm/common/basemodel/batch_objs.py | 8 +- lightllm/common/basemodel/cuda_graph.py | 14 +- .../linear_att/mtp_state_params.py | 49 +- .../server/router/model_infer/infer_batch.py | 2 - .../mode_backend/chunked_prefill/impl.py | 3 +- .../mode_backend/dp_backend/impl.py | 8 +- .../model_infer/mtp_speculative/utils.py | 10 +- test/att_resources/validate_model_cache.py | 12 +- test/benchmark/kernels/benchmark_kda_ssm.py | 14 +- .../att_resources/test_backend_startup.py | 13 +- .../att_resources/test_bounded_state.py | 3 +- .../test_hybrid_layer_mapping.py | 4 +- .../att_resources/test_mtp_state_commit.py | 58 +- .../basemodel/attention/linear/test_gdn.py | 42 +- .../basemodel/test_cuda_graph_autotune.py | 1 - .../linear_att/test_acceptance_fusion.py | 25 +- .../linear_att/test_mtp_state_params.py | 17 +- .../linear_att/test_replayssm.py | 13 +- .../linear_att/test_replayssm_execution.py | 15 +- .../linear_att/test_ssm_autotune_memory.py | 12 +- .../test_hybrid_req_manager_pd_sharing.py | 22 +- unit_tests/common/test_linear_state.py | 72 +- .../common/test_linear_state_contract.py | 149 +-- unit_tests/models/glm5_next/test_mtp.py | 8 +- unit_tests/models/test_post_layer_output.py | 7 +- 52 files changed, 3935 insertions(+), 2128 deletions(-) delete mode 100644 lightllm/common/attention_state/bounded_att/conv_ssm_state.py create mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py create mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/__init__.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/base.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/export.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/forward.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/native.py delete mode 100644 lightllm/common/attention_state/bounded_att/ssm/replay.py delete mode 100644 lightllm/common/basemodel/attention/linear/state.py diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py index 4efb909f0a..55190bea6d 100644 --- a/lightllm/common/attention_state/bounded_att/__init__.py +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -2,9 +2,9 @@ from .kda.create_kda_state import create_kda_state from .conv_ssm_types import ConvSSMStateConfig from .gdn.gdn_pages_pool import GDNPagesPool -from .gdn.gdn_state import GDNNativeState, GDNReplayState, GDNState +from .gdn.gdn_state import GDNState from .kda.kda_pages_pool import KDAPagesPool -from .kda.kda_state import KDANativeState, KDAReplayState, KDAState +from .kda.kda_state import KDAState from .base_state import BoundedAttentionState from .sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState @@ -16,13 +16,9 @@ "BoundedAttentionState", "ConvSSMStateConfig", "GDNState", - "GDNNativeState", - "GDNReplayState", "GDNPagesPool", "create_gdn_state", "KDAState", - "KDANativeState", - "KDAReplayState", "KDAPagesPool", "create_kda_state", "FP8E8M0SlidingWindowConfig", diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py deleted file mode 100644 index 62e3e83399..0000000000 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py +++ /dev/null @@ -1,590 +0,0 @@ -from __future__ import annotations - -from functools import lru_cache -from typing import TYPE_CHECKING, Any, NamedTuple -from math import prod - -import torch - -from lightllm.utils.log_utils import init_logger - -from ..commit import StateAcceptance -from .ssm import create_ssm_executor -from .ssm.forward import SSMForward -from .ssm.export import export_indexed_state - -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 - state_mode: str | None = None - - 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[str, object] | None = None, - gate_lower_bound: float = -5.0, - ) -> None: - """创建共享 Conv/SSM 运行态，并按模式绑定对应的 SSM executor。""" - self.config = config - if self.state_mode is not None and ssm_state_mode != self.state_mode: - raise ValueError(f"state mode {ssm_state_mode!r} does not match {self.state_mode!r}") - ssm_state_mode = self.state_mode or ssm_state_mode - 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_executor = create_ssm_executor( - 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) -> torch.Tensor: - """返回 executor 持有的 canonical SSM state。""" - return self._ssm_executor.state - - @property - def ssm_slots_per_req(self) -> int: - """返回每个请求占用的物理 SSM 槽位数。""" - return self._ssm_executor.slots_per_request - - def _reset_tracking(self, req_idx: int) -> None: - """清理请求的临时 history 和 MTP 接受游标。""" - self._ssm_executor.clear_history(req_idx) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].zero_() - - def reset(self, req_idx: int) -> None: - """清空请求的 Conv/SSM 运行态，同时保留页池快照。""" - 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: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """合并已接受 history，并返回 prefill 使用的请求和物理索引。""" - self._ssm_executor.merge_accepted_updates(reqs) - return reqs, reqs * self.ssm_slots_per_req - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> SSMForward: - """准备普通或 speculative decode，并返回对应的 SSMForward 上下文。""" - 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_executor.prepare_decode(reqs, cu_seqlens) - if not speculative: - self._ssm_executor.accept_single_token(reqs) - return SSMForward( - self, self._ssm_executor, reqs, self._ssm_executor.decode_indices(reqs, speculative), cu_seqlens, accepted - ) - - def get_accepted_state(self, req_idx: int) -> ConvSSMTensors: - """返回 canonical Conv/SSM 状态；Native 会选槽位，Replay 会物化快照。""" - 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_executor.get_accepted_state(req_idx, accepted)) - - def export_accepted_state(self, req_indices: list[int], targets: list[ConvSSMTensors]) -> None: - """在当前 stream 写入互不重叠的目标页，不分配整批临时快照。""" - export_indexed_state(self.conv_state, req_indices, [t.conv for t in targets], self.mtp_state_index) - self._ssm_executor.export_accepted_state(req_indices, [t.ssm for t in targets], self.mtp_state_index) - - def prepare_restore(self, req_idx: int) -> ConvSSMTensors: - """清理请求 tracking，并返回可写入恢复快照的 Conv/SSM 视图。""" - - 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: int, conv: torch.Tensor, ssm: torch.Tensor) -> None: - """把 Conv/SSM 快照复制回指定请求并完成恢复后的 tracking 初始化。""" - 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: int) -> int: - """把逻辑层号映射到当前 State 的物理层号。""" - return self._logical_to_storage_layer[logical_layer_id] - - def get_state_type(self) -> str: - """返回包含算法和 Conv/SSM 精度的稳定 State 类型名。""" - 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], - flush_history: bool = False, - ) -> 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) - # Hybrid radix page 是持久 checkpoint。ReplaySSM 可以在一个 kernel 中把 - # accepted ring 折叠到 canonical state 和目标页，使下一轮 decode 从空 history 开始。 - # MTP 使用独立的 Conv candidate 布局，因此保留原有行为。 - replay_flush = flush_history and self.mtp_state_index is None - targets = [ConvSSMTensors(pool.conv_state[slot], pool.ssm_state[slot]) for slot in page_slot_ids] - export_indexed_state(self.conv_state, req_indices, [target.conv for target in targets], self.mtp_state_index) - self._ssm_executor.export_accepted_state( - req_indices, - [target.ssm for target in targets], - self.mtp_state_index, - flush=replay_flush, - ) - - 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, acceptance: StateAcceptance) -> object | None: - """校验 acceptance 设备并委托 executor 检查可选的提交配置。""" - assert ( - acceptance.req_indices.device == self.ssm_state.device - ), "commit inputs must match the state device; CUDA states require CUDA tensors" - if self.mtp_state_index is not None: - validator = getattr(self._ssm_executor, "validate_commit", None) - return validator(acceptance, self.state_name) if validator is not None else None - return None - - def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: - """应用已接受的 MTP 前缀并推进对应的 executor 状态。""" - commit_data = self.validate_bounded_att_commit(acceptance) - if self.mtp_state_index is not None: - self._ssm_executor.commit_validated( - acceptance, self.mtp_state_index, self.state_name, commit_data=commit_data - ) - - # ======================================================================== - # 三、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" - remote_dp_bounded_state.export_accepted_state( - src_req_indices, [self.prepare_restore(req) for req in dst_req_indices] - ) - - # ======================================================================== - # 五、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) -> Any: - """根据大小页类型取得已启用的页池，不访问 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) -> dict[str, TensorField]: - """把 Conv/SSM 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/__init__.py b/lightllm/common/attention_state/bounded_att/gdn/__init__.py index 054e9bb519..62dc94c11d 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/__init__.py +++ b/lightllm/common/attention_state/bounded_att/gdn/__init__.py @@ -1,11 +1,5 @@ from .create_gdn_state import create_gdn_state from .gdn_pages_pool import GDNPagesPool -from .gdn_state import GDNNativeState, GDNReplayState, GDNState +from .gdn_state import GDNState -__all__ = [ - "create_gdn_state", - "GDNPagesPool", - "GDNState", - "GDNNativeState", - "GDNReplayState", -] +__all__ = ["create_gdn_state", "GDNPagesPool", "GDNState"] diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 0ba7d1f112..dcf92f65c4 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .gdn_state import GDNNativeState, GDNReplayState, GDNState +from .gdn_state import GDNState def create_gdn_state( @@ -15,22 +15,26 @@ def create_gdn_state( device: str | torch.device = "cuda", *, ssm_state_mode: str = "native", - ssm_replay_config: dict[str, object] | None = None, + ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, ) -> GDNState: - """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" - state_cls = { - "native": GDNNativeState, - "replay": GDNReplayState, - }.get(ssm_state_mode) - if state_cls is None: + """创建 GDN 请求运行态。replay 模式使用独立的 ReplayGDNState。""" + if ssm_state_mode == "replay": + from .replay_gdn_state import ReplayGDNState + + return ReplayGDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": raise ValueError(f"unknown GDN state mode: {ssm_state_mode}") - return state_cls( + return GDNState( config=config, 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 12cb426a7c..95cdef3697 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -1,23 +1,528 @@ -"""GDN 运行态及其 SSM 执行策略。""" +from __future__ import annotations -from ..conv_ssm_state import _ConvSSMState +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig from .gdn_pages_pool import GDNPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class GDNState(BoundedAttentionState): + """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" + + big_pages_pool: GDNPagesPool | None + small_pages_pool: GDNPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + @property + def ssm_slots_per_req(self) -> int: + """每个请求的物理 SSM 槽位数，包含 MTP 候选。""" + + return self.mtp_step + 1 + + 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) -class GDNState(_ConvSSMState): - """GDN 的公共资源外壳；具体 replay 策略由下列子类固定。""" + 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) - algorithm = "gdn" - pages_pool_class = GDNPagesPool + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + 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, + ) -class GDNNativeState(GDNState): - """GDN 原生 MTP 槽位策略。""" + # 阶段二：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 - state_mode = "native" + # 阶段三：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 -class GDNReplayState(GDNState): - """GDN replay history 策略。""" + 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)] + ) - state_mode = "replay" + @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/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py new file mode 100644 index 0000000000..5954b3846f --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -0,0 +1,1102 @@ +"""GDN replay 运行态。 + +这个文件单独实现 GDN 的 replay history、checkpoint 和 decode 准备。 +它不继承 native ``GDNState``，也不和 KDA replay 共享执行器基类。 +""" + +from __future__ import annotations + +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 collections.abc import Sequence + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix + + +def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +class ReplayGDNRuntime: + """以短 history 重放 decode 的 SSM 执行器。 + + ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 + cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode + 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 + checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 + """ + + _run_config: SSMRunConfig | None = None + + @property + def run_config(self) -> SSMRunConfig | None: + """返回最近一次 forward 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) + + def __init__( + self, + state: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = False, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> None: + """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 + + ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 + """ + + 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") + 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 + # 保留冗余布局，兼容已有调用方和 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) + # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 + # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 + # 下可以使用更窄的 history buffer；其他形状使用 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: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + 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, + 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, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor 以保留每个请求本轮已接受的前缀。 + + MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft + 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 + """ + _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: torch.Tensor) -> None: + """把接受的 history 物化回 canonical state，并清空请求 cursor。""" + 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: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + + layers, slots, hv, k, v = self.state.shape + config = self.run_config or SSMRunConfig(32, 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, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: + """返回不修改 active state/history 的单请求 canonical 快照。""" + 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 export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" + + if not req_indices: + return + shape = (self.state.shape[0], *self.state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) + self._materialize_accepted_state( + upload_indices(req_indices, self.state.device, torch.int32), + upload_indices([t.data_ptr() for t in targets], self.state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" + + return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward( + self, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, + *, + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" + + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + 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 == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor + # 运算，无法定义静态 shape。 + 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, + self.kda, + self.gate_lower_bound, + 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: torch.Tensor) -> None: + """非 MTP decode 直接把单 token 写入当前 history。""" + self.accept_updates(reqs) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.state.shape[0], *self.state.shape[2:]), + dtype=self.state.dtype, + device=self.state.device, + ) + self._materialize_accepted_state( + torch.tensor([req], dtype=torch.int32, device=self.state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.state[:, req].copy_(snap) + self.clear_history(req) + + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from .gdn_pages_pool import GDNPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class ReplayGDNState(BoundedAttentionState): + """GDN replay：每个请求一份 canonical SSM，外加一段不进入 checkpoint 的 history。""" + + 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", + *, + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """按请求容量创建 Conv 运行态、单槽 canonical SSM 和 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayGDNState only implements replay mode, got {ssm_state_mode}") + 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.gate_lower_bound = gate_lower_bound + 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, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) + self.runtime = ReplayGDNRuntime( + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + num_key_heads=config.num_k_heads, + kda=False, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + ) + 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 + + @property + def ssm_slots_per_req(self) -> int: + """Replay 每个请求只保留一个 canonical SSM 槽位。""" + + return 1 + + @property + def _ssm_executor(self): + """供既有测试读取 history；实现本身不依赖这个名字。""" + + return self.runtime + + 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_() + self.ssm_state[:, req_idx].zero_() + self.runtime.clear_history(req_idx) + 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) + conv_copies = [ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + copy_bounded_att_tensors(copies=conv_copies) + self.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；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) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.runtime.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + 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: + """提交已接受前缀：更新 MTP 游标，并把同样的长度加到 replay cursor。""" + if self.mtp_state_index is None: + return + self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """Prefill 前把已接受 history 折回 canonical state。""" + + self.runtime.merge_accepted_updates(reqs) + + def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: + """折叠即将溢出的 history。非 MTP decode 同时接受本轮唯一 token。""" + + self.runtime.prepare_decode(reqs, cu_seqlens) + if cu_seqlens is None: + self.runtime.accept_single_token(reqs) + + def replay_forward( + self, + logical_layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> torch.Tensor: + """执行一层 GDN replay decode。""" + + output, _ = self.runtime.forward( + self._logical_to_storage_layer[logical_layer], + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens, + ) + return output + + def get_accepted_state(self, req_idx: int): + """返回已接受 Conv 窗口和物化后的 SSM，不改写 canonical state 或 cursor。""" + + ssm = self.runtime.get_accepted_state(req_idx) + return _AcceptedTensors(self._accepted_conv(req_idx), ssm) + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + index = torch.arange(width, device=conv.device) + offset + return conv.index_select(-1, index) + + # ======================================================================== + # 三、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: ReplayGDNState, + ) -> 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: ReplayGDNState, + ) -> 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" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.runtime.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.runtime.export_ssm( + src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices] + ) + 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: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._fields(accepted.conv, accepted.ssm) + 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" + self.runtime.clear_history(req_idx) + fields = self._fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + 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]) + + @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")), + } + + +from typing import NamedTuple + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode + + +def _find_replay_state(att_state): + manager = getattr(att_state.infer_state, "att_state_manager", None) + if manager is None: + return None + role = att_state.infer_state.att_state_role + matches = [state for state in manager.get_bounded_attention_states(role) if isinstance(state, ReplayGDNState)] + if not matches: + return None + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}, found {len(matches)}") + return matches[0] + + +def prepare_replay_gdn_prefill(att_state) -> bool: + """Replay prefill 先折叠 history，并用单槽 SSM 索引。""" + + state = _find_replay_state(att_state) + if state is None: + return False + reqs = att_state.infer_state.b_req_idx + state.merge_accepted_updates(reqs) + att_state.b_conv_buffer_idx = reqs + att_state.b_ssm_buffer_idx = reqs + att_state._replay_state = state + return True + + +def prepare_replay_gdn_decode(att_state) -> bool: + """准备 GDN replay decode 的 conv/MTP 索引，并折叠本轮 history。""" + + state = _find_replay_state(att_state) + if state is None: + return False + draft_step = att_state.backend.model.mtp_manager.get_decode_draft_step(att_state.backend.model.is_mtp_draft_model) + infer = att_state.infer_state + if draft_step == 0: + att_state.b_conv_buffer_idx = infer.b_req_idx + att_state.b_ssm_buffer_idx = infer.b_req_idx + att_state.b1_mtp_cu_q_seq_len = None + att_state.b_num_accepted_tokens = None + else: + if att_state.backend.uses_dynamic_spec_verify_layout(): + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, + ) + + ( + att_state.b1_mtp_cu_q_seq_len, + att_state.b_conv_buffer_idx, + att_state.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=state.mtp_state_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + else: + width = draft_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "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_state.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + att_state.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + att_state.b_num_accepted_tokens = state.mtp_state_index[att_state.b_conv_buffer_idx] + 1 + att_state.b_ssm_buffer_idx = att_state.b_conv_buffer_idx + state.prepare_for_decode(att_state.b_conv_buffer_idx, att_state.b1_mtp_cu_q_seq_len) + att_state._replay_state = state + return True diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index 468fd0da28..fe1fceb3bd 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -4,12 +4,10 @@ from abc import ABC, abstractmethod from functools import lru_cache -from typing import Any +from typing import Any, List import torch -from ..commit import StateAcceptance - from ..utils import deepcopy_result from lightllm.common.attention_state.types import PageKind, StateTransferLayout, TransferRegion @@ -65,34 +63,32 @@ def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod - def reset_bounded_att_state(self, req_indices: list[int]) -> None: + def reset_bounded_att_state(self, req_indices: List[int]) -> None: """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" ... @abstractmethod def save_bounded_att_runtime_to_pages( self, - req_indices: list[int], - state_token_lens: list[int], + req_indices: List[int], + state_token_lens: List[int], page_kind: PageKind, - page_slot_ids: list[int], - flush_history: bool = False, + page_slot_ids: List[int], ) -> None: """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 - ``flush_history`` 用于持久化 radix checkpoint：支持的 State 可在写页时 - 同步折叠临时 history；前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 + 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 """ ... @abstractmethod def restore_bounded_att_runtime_from_pages( self, - req_indices: list[int], - state_token_lens: list[int], + req_indices: List[int], + state_token_lens: List[int], page_kind: PageKind, - page_slot_ids: list[int], + page_slot_ids: List[int], ) -> None: """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 @@ -102,7 +98,12 @@ def restore_bounded_att_runtime_from_pages( ... @abstractmethod - def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" ... @@ -129,9 +130,9 @@ def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: def write_state_pages_to_cpu( self, page_kind: PageKind, - page_slot_ids: list[int], + page_slot_ids: List[int], cpu_page_buffers: torch.Tensor, - cpu_page_indices: list[int], + cpu_page_indices: List[int], layout: StateTransferLayout, cpu_page_ready_flags: list[bool] | None = None, ) -> None: @@ -142,8 +143,8 @@ def write_state_pages_to_cpu( def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, - cpu_page_indices: list[int], - big_page_slot_ids: list[int], + cpu_page_indices: List[int], + big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" @@ -158,8 +159,8 @@ class DpCopyInterface(ABC): @abstractmethod def copy_big_pages_from_dp( self, - src_page_ids: list[int], - dst_page_ids: list[int], + src_page_ids: List[int], + dst_page_ids: List[int], remote_dp_bounded_state: CoreInterface, ) -> None: """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" @@ -168,11 +169,14 @@ def copy_big_pages_from_dp( @abstractmethod def copy_runtime_state_from_dp( self, - src_req_indices: list[int], - dst_req_indices: list[int], + src_req_indices: List[int], + dst_req_indices: List[int], remote_dp_bounded_state: CoreInterface, ) -> None: + """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ ... @@ -201,7 +205,10 @@ 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/__init__.py b/lightllm/common/attention_state/bounded_att/kda/__init__.py index 498ba3b227..7154892202 100644 --- a/lightllm/common/attention_state/bounded_att/kda/__init__.py +++ b/lightllm/common/attention_state/bounded_att/kda/__init__.py @@ -1,11 +1,5 @@ from .create_kda_state import create_kda_state from .kda_pages_pool import KDAPagesPool -from .kda_state import KDANativeState, KDAReplayState, KDAState +from .kda_state import KDAState -__all__ = [ - "create_kda_state", - "KDAPagesPool", - "KDAState", - "KDANativeState", - "KDAReplayState", -] +__all__ = ["create_kda_state", "KDAPagesPool", "KDAState"] diff --git a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py index af90adad00..f94b57501d 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 @@ -5,7 +5,7 @@ import torch from ..conv_ssm_types import ConvSSMStateConfig -from .kda_state import KDANativeState, KDAReplayState, KDAState +from .kda_state import KDAState def create_kda_state( @@ -15,22 +15,26 @@ def create_kda_state( device: str | torch.device = "cuda", *, ssm_state_mode: str = "native", - ssm_replay_config: dict[str, object] | None = None, + ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, ) -> KDAState: - """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" - state_cls = { - "native": KDANativeState, - "replay": KDAReplayState, - }.get(ssm_state_mode) - if state_cls is None: + """创建 KDA 请求运行态。replay 模式使用独立的 ReplayKDAState。""" + if ssm_state_mode == "replay": + from .replay_kda_state import ReplayKDAState + + return ReplayKDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": raise ValueError(f"unknown KDA state mode: {ssm_state_mode}") - return state_cls( + return KDAState( config=config, 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 6723e8aeec..5bb9fc780d 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -1,23 +1,528 @@ -"""KDA 运行态及其 SSM 执行策略。""" +from __future__ import annotations -from ..conv_ssm_state import _ConvSSMState +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig from .kda_pages_pool import KDAPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class KDAState(BoundedAttentionState): + """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" + + big_pages_pool: KDAPagesPool | None + small_pages_pool: KDAPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + @property + def ssm_slots_per_req(self) -> int: + """每个请求的物理 SSM 槽位数，包含 MTP 候选。""" + + return self.mtp_step + 1 + + 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) -class KDAState(_ConvSSMState): - """KDA 的公共资源外壳；具体 replay 策略由下列子类固定。""" + 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) - algorithm = "kda" - pages_pool_class = KDAPagesPool + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + 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, + ) -class KDANativeState(KDAState): - """KDA 原生 MTP 槽位策略。""" + # 阶段二：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 - state_mode = "native" + # 阶段三：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 -class KDAReplayState(KDAState): - """KDA replay history 策略。""" + 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)] + ) - state_mode = "replay" + @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/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py new file mode 100644 index 0000000000..2b4e33697c --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -0,0 +1,1102 @@ +"""KDA replay 运行态。 + +这个文件单独实现 KDA 的 replay history、checkpoint 和 decode 准备。 +它不继承 native ``GDNState``，也不和 GDN replay 共享执行器基类。 +""" + +from __future__ import annotations + +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 collections.abc import Sequence + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix + + +def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +class ReplayKDARuntime: + """以短 history 重放 decode 的 SSM 执行器。 + + ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 + cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode + 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 + checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 + """ + + _run_config: SSMRunConfig | None = None + + @property + def run_config(self) -> SSMRunConfig | None: + """返回最近一次 forward 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) + + def __init__( + self, + state: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = True, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> None: + """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 + + ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 + """ + + 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") + 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 + # 保留冗余布局，兼容已有调用方和 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) + # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 + # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 + # 下可以使用更窄的 history buffer；其他形状使用 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: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + 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, + 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, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor 以保留每个请求本轮已接受的前缀。 + + MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft + 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 + """ + _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: torch.Tensor) -> None: + """把接受的 history 物化回 canonical state，并清空请求 cursor。""" + 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: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + + layers, slots, hv, k, v = self.state.shape + config = self.run_config or SSMRunConfig(32, 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, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: + """返回不修改 active state/history 的单请求 canonical 快照。""" + 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 export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" + + if not req_indices: + return + shape = (self.state.shape[0], *self.state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) + self._materialize_accepted_state( + upload_indices(req_indices, self.state.device, torch.int32), + upload_indices([t.data_ptr() for t in targets], self.state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" + + return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward( + self, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, + *, + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" + + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + 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 == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor + # 运算，无法定义静态 shape。 + 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, + self.kda, + self.gate_lower_bound, + 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: torch.Tensor) -> None: + """非 MTP decode 直接把单 token 写入当前 history。""" + self.accept_updates(reqs) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.state.shape[0], *self.state.shape[2:]), + dtype=self.state.dtype, + device=self.state.device, + ) + self._materialize_accepted_state( + torch.tensor([req], dtype=torch.int32, device=self.state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.state[:, req].copy_(snap) + self.clear_history(req) + + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from .kda_pages_pool import KDAPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class ReplayKDAState(BoundedAttentionState): + """KDA replay：每个请求一份 canonical SSM，外加一段不进入 checkpoint 的 history。""" + + 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", + *, + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """按请求容量创建 Conv 运行态、单槽 canonical SSM 和 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayKDAState only implements replay mode, got {ssm_state_mode}") + 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.gate_lower_bound = gate_lower_bound + 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, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) + self.runtime = ReplayKDARuntime( + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + num_key_heads=config.num_k_heads, + kda=True, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + ) + 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 + + @property + def ssm_slots_per_req(self) -> int: + """Replay 每个请求只保留一个 canonical SSM 槽位。""" + + return 1 + + @property + def _ssm_executor(self): + """供既有测试读取 history；实现本身不依赖这个名字。""" + + return self.runtime + + def get_state_type(self) -> str: + """返回 GDN 算法及 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_() + self.ssm_state[:, req_idx].zero_() + self.runtime.clear_history(req_idx) + 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) + conv_copies = [ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + copy_bounded_att_tensors(copies=conv_copies) + self.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；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) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.runtime.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + 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: + """提交已接受前缀：更新 MTP 游标，并把同样的长度加到 replay cursor。""" + if self.mtp_state_index is None: + return + self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """Prefill 前把已接受 history 折回 canonical state。""" + + self.runtime.merge_accepted_updates(reqs) + + def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: + """折叠即将溢出的 history。非 MTP decode 同时接受本轮唯一 token。""" + + self.runtime.prepare_decode(reqs, cu_seqlens) + if cu_seqlens is None: + self.runtime.accept_single_token(reqs) + + def replay_forward( + self, + logical_layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> torch.Tensor: + """执行一层 KDA replay decode。""" + + output, _ = self.runtime.forward( + self._logical_to_storage_layer[logical_layer], + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens, + ) + return output + + def get_accepted_state(self, req_idx: int): + """返回已接受 Conv 窗口和物化后的 SSM，不改写 canonical state 或 cursor。""" + + ssm = self.runtime.get_accepted_state(req_idx) + return _AcceptedTensors(self._accepted_conv(req_idx), ssm) + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + index = torch.arange(width, device=conv.device) + offset + return conv.index_select(-1, index) + + # ======================================================================== + # 三、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: ReplayKDAState, + ) -> 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: ReplayKDAState, + ) -> 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" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.runtime.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.runtime.export_ssm( + src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices] + ) + 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: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._fields(accepted.conv, accepted.ssm) + 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" + self.runtime.clear_history(req_idx) + fields = self._fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + 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]) + + @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")), + } + + +from typing import NamedTuple + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode + + +def _find_replay_state(att_state): + manager = getattr(att_state.infer_state, "att_state_manager", None) + if manager is None: + return None + role = att_state.infer_state.att_state_role + matches = [state for state in manager.get_bounded_attention_states(role) if isinstance(state, ReplayKDAState)] + if not matches: + return None + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}, found {len(matches)}") + return matches[0] + + +def prepare_replay_kda_prefill(att_state) -> bool: + """Replay prefill 先折叠 history，并用单槽 SSM 索引。""" + + state = _find_replay_state(att_state) + if state is None: + return False + reqs = att_state.infer_state.b_req_idx + state.merge_accepted_updates(reqs) + att_state.b_conv_buffer_idx = reqs + att_state.b_ssm_buffer_idx = reqs + att_state._replay_state = state + return True + + +def prepare_replay_kda_decode(att_state) -> bool: + """准备 KDA replay decode 的 conv/MTP 索引，并折叠本轮 history。""" + + state = _find_replay_state(att_state) + if state is None: + return False + draft_step = att_state.backend.mtp_step + infer = att_state.infer_state + if draft_step == 0: + att_state.b_conv_buffer_idx = infer.b_req_idx + att_state.b_ssm_buffer_idx = infer.b_req_idx + att_state.b1_mtp_cu_q_seq_len = None + att_state.b_num_accepted_tokens = None + else: + if att_state.backend.uses_dynamic_spec_verify_layout(): + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, + ) + + ( + att_state.b1_mtp_cu_q_seq_len, + att_state.b_conv_buffer_idx, + att_state.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=state.mtp_state_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + else: + width = draft_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "KDA fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + att_state.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + att_state.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + att_state.b_num_accepted_tokens = state.mtp_state_index[att_state.b_conv_buffer_idx] + 1 + att_state.b_ssm_buffer_idx = att_state.b_conv_buffer_idx + state.prepare_for_decode(att_state.b_conv_buffer_idx, att_state.b1_mtp_cu_q_seq_len) + att_state._replay_state = state + return True 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 d1455cdb34..a1d72102ac 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 @@ -9,7 +9,6 @@ import torch -from ...commit import StateAcceptance from ...utils import deepcopy_result from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState @@ -154,7 +153,6 @@ def save_bounded_att_runtime_to_pages( state_token_lens: list[int], page_kind: PageKind, page_slot_ids: list[int], - flush_history: bool = False, ) -> None: pool = self._pool(page_kind=page_kind) for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): @@ -181,7 +179,7 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> 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 deleted file mode 100644 index 3cf73492ec..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from ..conv_ssm_types import ConvSSMStateConfig - from .base import SSMExecutor - - -def create_ssm_executor( - config: ConvSSMStateConfig, - request_capacity: int, - mtp_step: int, - device: str | torch.device, - mode: str, - replay_config: dict[str, object] | None, - *, - kda: bool, - gate_lower_bound: float, -) -> SSMExecutor: - """根据状态模式创建 Native 或 Replay SSM executor。""" - """按统一 state layout 创建 native 或 replay executor。 - - native 为每个 MTP 位置分配物理槽位，replay 为每个请求分配一个槽位， - 并在 executor 内保存短 history。 - """ - - if mode not in ("native", "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 NativeSSMExecutor - - return NativeSSMExecutor(state, mtp_step + 1, **kwargs) - kwargs.update(num_key_heads=config.num_k_heads) - from .replay import ReplaySSMExecutor - - replay_config = replay_config or {} - return ReplaySSMExecutor( - state, - replay_config.get("cache_len", 16), - mtp_step + 1, - config.conv_dtype, - compute_mode=replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline"), - **kwargs, - ) diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py deleted file mode 100644 index 68c2adcc7a..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/base.py +++ /dev/null @@ -1,130 +0,0 @@ -from __future__ import annotations - -import torch - -from lightllm.common.attention_state.commit import StateAcceptance -from lightllm.common.ssm_types import SSMRunConfig - -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from .export import export_indexed_state - - -def update_accepted_indices( - acceptance: StateAcceptance, indices: torch.Tensor, *, replay_cursors: torch.Tensor | None = None -) -> None: - """把接受长度写入 MTP 槽位索引，并按需推进 Replay cursor。""" - - commit_ssm_prefix(indices, acceptance.req_indices, acceptance.accepted_lengths, replay_cursors) - - -class SSMExecutor: - """一种 SSM 模式的 decode、接受和快照适配器。 - - 算法 State 持有资源和逻辑层到物理层的映射；executor 只处理 decode 递推、接受前缀 - 以及 accepted-state 导出。prefill 仍由 attention backend 执行。一次 verify 的 - 布局由 ``SSMForward`` 保存，并在 commit 阶段通过 ``StateCommit`` 传回同一个 executor。 - """ - - # State 负责资源生命周期和层映射，Executor 只负责本模式的 decode、接受和快照导出。 - # 这样 Native、Replay 共享同一条上层调用链，差异收敛在执行器内部。 - - slots_per_request = 1 - _run_config: SSMRunConfig | None = None - - @property - def run_config(self) -> SSMRunConfig | None: - """返回最近一次 forward 使用的不可变 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: SSMRunConfig | dict[str, int] | None) -> None: - """保存不可变的 kernel 配置，兼容字典形式的 autotune 结果。""" - self._run_config = config if config is None or isinstance(config, SSMRunConfig) else SSMRunConfig(**config) - - def clear_history(self, req: int | torch.Tensor) -> None: - """清空请求的临时 history；默认模式没有额外 history。""" - - return None - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> None: - """在 decode 前准备本轮基准；Native 默认不需要准备。""" - - return None - - def accept_single_token(self, reqs: torch.Tensor) -> None: - """提交非 speculative decode 产生的单 token 更新。""" - - return None - - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把已接受的临时更新合并回 canonical state。""" - - return None - - def decode_indices(self, reqs: torch.Tensor, speculative: bool) -> torch.Tensor: - """返回 decode 使用的物理槽位；Native MTP 会展开请求槽位。""" - - return reqs - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回一个请求的 accepted-state 视图或物化快照。""" - - return self.state[:, req_idx] - - def export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """把 accepted state 写入调用方提供的目标页，不负责分配目标。""" - - export_indexed_state(self.state, req_indices, targets) - - def create_commit(self, run_config: SSMRunConfig | None) -> object | None: - """把本次 forward 的配置转换为接受阶段可携带的提交 payload。""" - - return None - - def commit(self, acceptance: StateAcceptance, indices: torch.Tensor, state_name: str) -> None: - """校验并应用接受结果；供直接调用 State 接口时使用。""" - - validator = getattr(self, "validate_commit", None) - commit_data = validator(acceptance, state_name) if validator is not None else None - self.commit_validated(acceptance, indices, state_name, commit_data=commit_data) - - def commit_validated( - self, - acceptance: StateAcceptance, - indices: torch.Tensor, - state_name: str, - *, - commit_data: object | None = None, - ) -> None: - """应用已由 State/Manager 校验过的接受结果，不重复检查公共契约。""" - - update_accepted_indices(acceptance, indices) - - def forward( - self, - layer: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig | None]: - """执行一层 decode，并返回输出及本次实际采用的 run config。""" - - raise NotImplementedError diff --git a/lightllm/common/attention_state/bounded_att/ssm/export.py b/lightllm/common/attention_state/bounded_att/ssm/export.py deleted file mode 100644 index bfb55bc535..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/export.py +++ /dev/null @@ -1,104 +0,0 @@ -from collections.abc import Sequence -from math import prod - -import torch -import triton -import triton.language as tl - -from ..copy_kernel import copy_bounded_att_tensors - - -def upload_indices(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device，使用 pinned memory 减少同步。""" - - return torch.tensor(values, dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -@triton.jit(do_not_specialize=["count"]) -def _export_indexed( - Source, - Targets, - Reqs, - Accepted, - count, - SHAPE: tl.constexpr, - SOURCE_STRIDES: tl.constexpr, - TARGET_STRIDES: tl.constexpr, - REQUEST_STRIDE: tl.constexpr, - INDEX_STRIDE: tl.constexpr, - SIZE: tl.constexpr, - BLOCK: tl.constexpr, -): - # 每个 request 对应一个预分配目标指针；kernel 只做 indexed copy， - # 不改变 source cursor/history，也不负责目标页生命周期。 - total = count.to(tl.int64) * SIZE - for start in range(tl.program_id(0).to(tl.int64) * BLOCK, total, tl.num_programs(0) * BLOCK): - row = start + tl.arange(0, BLOCK) - valid = row < total - item, offset = row // SIZE, row % SIZE - req = tl.load(Reqs + item, valid, 0).to(tl.int64) - accepted = tl.load(Accepted + req, valid, 0).to(tl.int64) - source_offset = req * REQUEST_STRIDE + accepted * INDEX_STRIDE - target_offset = tl.full((BLOCK,), 0, tl.int64) - for axis in tl.static_range(len(SHAPE) - 1, -1, -1): - position = offset % SHAPE[axis] - offset //= SHAPE[axis] - source_offset += position * SOURCE_STRIDES[axis] - target_offset += position * TARGET_STRIDES[axis] - target = tl.load(Targets + item, valid, 0).to(tl.pointer_type(Source.dtype.element_ty)) - values = tl.load(Source + source_offset, valid, 0) - tl.store(target + target_offset, values, valid) - - -def export_indexed_state( - source: torch.Tensor, - req_indices: Sequence[int], - targets: Sequence[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - slots_per_request: int = 1, - index_stride: int = 1, -) -> None: - """将 source 中的 accepted state 写入预分配的 target 页。 - - ``req_indices`` 和 ``targets`` 必须一一对应；MTP 通过 - ``accepted_indices`` 选择请求内槽位。该函数只负责复制，不拥有 target - 的分配和释放生命周期。 - """ - - if not req_indices: - return - assert len(req_indices) == len(targets) - shape, target_strides = targets[0].shape, targets[0].stride() - assert all(t.shape == shape and t.stride() == target_strides and t.dtype == source.dtype for t in targets) - if accepted_indices is None: - copies = [(source[:, req * slots_per_request][..., : shape[-1]], dst) for req, dst in zip(req_indices, targets)] - if source.is_cuda: - copy_bounded_att_tensors(copies) - else: - for src, dst in copies: - dst.copy_(src) - return - if not source.is_cuda: - for req, dst in zip(req_indices, targets): - accepted = int(accepted_indices[req]) - if slots_per_request > 1: - dst.copy_(source[:, req * slots_per_request + accepted]) - else: - dst.copy_(source[:, req, ..., accepted : accepted + shape[-1]]) - return - _export_indexed[(24,)]( - source, - upload_indices([t.data_ptr() for t in targets], source.device), - upload_indices(req_indices, source.device, torch.int32), - accepted_indices, - len(req_indices), - tuple(shape), - (source.stride(0), *source.stride()[2:]), - target_strides, - source.stride(1) * slots_per_request, - index_stride, - prod(shape), - 1024, - num_warps=8, - ) diff --git a/lightllm/common/attention_state/bounded_att/ssm/forward.py b/lightllm/common/attention_state/bounded_att/ssm/forward.py deleted file mode 100644 index a23054299a..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/forward.py +++ /dev/null @@ -1,98 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention_state.commit import StateCommit -from lightllm.common.ssm_types import SSMRunConfig - -if TYPE_CHECKING: - from .base import SSMExecutor - from ..conv_ssm_state import _ConvSSMState - - -@dataclass -class SSMForward: - """携带一次 forward 的请求布局，并生成对应的状态提交凭据。 - - ``req_indices`` 是本次 forward 的 request-row 序列，可能包含 CUDA Graph - 的 ``HOLD`` padding；``state_indices`` 是 kernel 使用的物理槽位， - ``cu_seqlens`` 描述真实 query 窗口。``run_config`` 在第一层执行时确定， - 后续层复用同一配置，``finish`` 再把它封装进 ``StateCommit``。 - """ - - resource: _ConvSSMState - executor: SSMExecutor - req_indices: torch.Tensor - state_indices: torch.Tensor - cu_seqlens: torch.Tensor | None - accepted_token_counts: torch.Tensor | None - run_config: SSMRunConfig | None = None - _layout_checked: bool = field(default=False, init=False, repr=False) - - def run_layer( - self, - logical_layer: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - ) -> torch.Tensor: - """执行一层线性注意力；logical_layer 会映射到物理权重层。""" - - output, self.run_config = self.executor.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) -> StateCommit | None: - """结束 forward，返回接受阶段需要的状态提交描述。""" - - payload = self.executor.create_commit(self.run_config) - if payload is None: - return None - return StateCommit(self.resource.state_name, payload) - - def copy_inputs(self, source: SSMForward) -> None: - """把动态图输入复制到已 capture 的 tensor。 - - graph bucket 的资源和 tensor layout 在首次绑定时校验一次，后续调用 - 只更新数据；captured layout 不随请求变化。 - """ - - names = ("req_indices", "state_indices", "cu_seqlens", "accepted_token_counts") - if not self._layout_checked: - if self.resource is not source.resource or self.executor is not source.executor: - raise ValueError("cannot copy a forward for a different resource or executor") - for name in names: - target, value = getattr(self, name), getattr(source, name) - if (target is None) != (value is None) or ( - target is not None - and (target.shape != value.shape or target.dtype != value.dtype or target.device != value.device) - ): - raise ValueError(f"captured SSM layout changed: {name}") - self._layout_checked = True - - copied = set() - for name in names: - target, value = getattr(self, name), getattr(source, name) - 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 deleted file mode 100644 index b664f78027..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/native.py +++ /dev/null @@ -1,130 +0,0 @@ -from __future__ import annotations - -import torch - -from lightllm.common.ssm_types import SSMRunConfig -from .base import SSMExecutor -from .export import export_indexed_state - - -class NativeSSMExecutor(SSMExecutor): - """直接调用 fused recurrent kernel，并把 MTP 状态展开为物理槽位。""" - - def __init__( - self, state: torch.Tensor, verify_width: int, *, kda: bool = False, gate_lower_bound: float = -5.0 - ) -> None: - """创建直接写入 canonical state 的 Native executor。""" - self.state = state - self.slots_per_request = verify_width - self.kda = kda - self.gate_lower_bound = gate_lower_bound - - def decode_indices(self, reqs: torch.Tensor, speculative: bool) -> torch.Tensor: - """返回普通 decode 或 speculative verify 对应的物理 state 索引。""" - - 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: int, accepted_index: int | None = None) -> torch.Tensor: - """返回单槽位 view，或 MTP 接受位置对应的 state 副本。""" - - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """把 native state 按请求和接受位置导出到调用方目标页。""" - - export_indexed_state( - self.state, - req_indices, - targets, - accepted_indices if self.slots_per_request > 1 else None, - slots_per_request=self.slots_per_request, - index_stride=self.state.stride(1), - ) - - def forward( - self, - layer: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, None]: - """执行 native recurrent decode，并原地更新 fused kernel 的最终 state。""" - - 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 deleted file mode 100644 index 51987ba61a..0000000000 --- a/lightllm/common/attention_state/bounded_att/ssm/replay.py +++ /dev/null @@ -1,359 +0,0 @@ -from __future__ import annotations - -import torch -import triton - -from lightllm.common.attention_state.commit import StateAcceptance -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 SSMExecutor, update_accepted_indices -from .export import upload_indices - - -class ReplaySSMExecutor(SSMExecutor): - """以短 history 重放 decode 的 SSM 执行器。 - - ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 - cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode - 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 - checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 - """ - - def __init__( - self, - state: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - *, - num_key_heads: int | None = None, - kda: bool = False, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, - ) -> None: - """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 - - ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 - """ - - 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") - 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 - # 保留冗余布局，兼容已有调用方和 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) - # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 - # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 - # 下可以使用更窄的 history buffer；其他形状使用 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: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - 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, - 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, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor 以保留每个请求本轮已接受的前缀。 - - MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft - 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 - """ - _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: torch.Tensor) -> None: - """把接受的 history 物化回 canonical state，并清空请求 cursor。""" - 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: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - - layers, slots, hv, k, v = self.state.shape - config = self.run_config or SSMRunConfig(32, 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, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回不修改 active state/history 的单请求 canonical 快照。""" - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" - - if not req_indices: - return - shape = (self.state.shape[0], *self.state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) - self._materialize_accepted_state( - upload_indices(req_indices, self.state.device, torch.int32), - upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> SSMRunConfig: - """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" - - return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig]: - """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" - - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - 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 == "precompute" - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor - # 运算，无法定义静态 shape。 - 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, - self.kda, - self.gate_lower_bound, - 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: torch.Tensor) -> None: - """非 MTP decode 直接把单 token 写入当前 history。""" - self.accept_updates(reqs) - - def commit(self, acceptance: StateAcceptance, indices: torch.Tensor, state_name: str) -> None: - """应用接受结果；MTP 直接推进 cursor，单 token 还需写入 history。""" - - self.commit_validated(acceptance, indices, state_name) - - def commit_validated( - self, - acceptance: StateAcceptance, - indices: torch.Tensor, - state_name: str, - *, - commit_data: object | None = None, - ) -> None: - """应用已验证的接受结果，推进 replay cursor。""" - - if self.verify_width > 1: - update_accepted_indices(acceptance, indices, replay_cursors=self.cursors) - else: - update_accepted_indices(acceptance, indices) - self.accept_updates(acceptance.req_indices, 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 54d92a2e4e..cea0a49af2 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 @@ -12,7 +12,6 @@ from lightllm.utils.envs_utils import get_env_start_args -from ...commit import StateAcceptance from ...utils import deepcopy_result from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState @@ -174,7 +173,6 @@ def save_bounded_att_runtime_to_pages( state_token_lens: list[int], page_kind: PageKind, page_slot_ids: list[int], - flush_history: bool = False, ) -> None: """边界必须完成一个产出步长；仅保留下一步仍需读取的窗口重叠部分。""" assert all( @@ -212,7 +210,7 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, acceptance: StateAcceptance) -> None: + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index bccb3c63e6..6fb18c7ea1 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -2,9 +2,9 @@ from typing import TYPE_CHECKING from functools import lru_cache -import torch +from typing import List -from .commit import StateAcceptance +import torch from lightllm.common.attention_state.types import PageKind @@ -331,7 +331,7 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # ======================================================================== # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 - def reset_bounded_att_state(self, req_indices: list[int]) -> None: + def reset_bounded_att_state(self, req_indices: List[int]) -> None: """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" for state in self.get_bounded_attention_states(): state.reset_bounded_att_state(req_indices) @@ -340,11 +340,10 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 def save_bounded_att_runtime_to_pages( self, - req_indices: list[int], - state_token_lens: list[int], + req_indices: List[int], + state_token_lens: List[int], page_kind: PageKind, - page_slot_ids: list[int], - flush_history: bool = False, + page_slot_ids: List[int], ) -> None: """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" bounded_states = self.get_bounded_attention_states() @@ -363,15 +362,14 @@ def save_bounded_att_runtime_to_pages( state_token_lens=state_token_lens, page_kind=page_kind, page_slot_ids=page_slot_ids, - flush_history=flush_history, ) def restore_bounded_att_runtime_from_pages( self, - req_indices: list[int], - state_token_lens: list[int], + req_indices: List[int], + state_token_lens: List[int], page_kind: PageKind, - page_slot_ids: list[int], + page_slot_ids: List[int], ) -> None: """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" bounded_states = self.get_bounded_attention_states() @@ -394,14 +392,17 @@ def restore_bounded_att_runtime_from_pages( ) # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 - def commit_bounded_att_state(self, acceptance: StateAcceptance, role: str) -> None: - """提交按请求排列的已接受前缀；无参与 State 时为空操作。""" - assert acceptance.req_indices.is_cuda, "commit inputs must be CUDA tensors" + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + role: str, + ) -> None: + """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" assert role in MODEL_ROLES, f"unknown model role {role}" - states = self.get_bounded_attention_states(role) - acceptance.validate({state.state_name for state in states}) - for state in states: - state.commit_bounded_att_state(acceptance) + for state in self.get_bounded_attention_states(role): + state.commit_bounded_att_state(req_indices, accepted_steps) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index b3aa1c8592..ae37385e68 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -10,7 +10,6 @@ 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 @@ -164,16 +163,6 @@ 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 79ad4e20e6..6e57eca732 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 GDNLinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend -class FlashQlaLinearAttBackend(GDNLinearAttBackend): +class FlashQlaLinearAttBackend(LinearAttBackend): 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 0e44139037..07d258bd24 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,17 +1,23 @@ 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 - -import torch - -from ..base_att import BaseAttBackend, AttControl +from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, 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 @@ -20,7 +26,7 @@ from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight -class GDNLinearAttBackend(BaseAttBackend, ABC): +class LinearAttBackend(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_) @@ -94,16 +100,33 @@ 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") -> "GDNLinearAttPrefillAttState": - return GDNLinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + 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_decode_state(self, infer_state: "InferStateInfo") -> "GDNLinearAttDecodeAttState": - return GDNLinearAttDecodeAttState(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)) @dataclasses.dataclass -class GDNLinearAttPrefillAttState(LinearPrefillAttState): - state_class = GDNState +class LinearAttPrefillAttState(BasePrefillAttState): + + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + self._replay_state = None + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import prepare_replay_gdn_prefill + + if prepare_replay_gdn_prefill(self): + return + 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 def prefill_att( self, @@ -120,24 +143,22 @@ 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: GDNLinearAttBackend = self.backend + backend: LinearAttBackend = self.backend - resource = get_resource(self, layer_num) - conv_states, ssm_states = resource.get_layer_view(layer_num) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=type(self._replay_state) if self._replay_state is not None else GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=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, - self.state_indices[resource.state_name], + mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight ) return core_attn_out, z @@ -150,19 +171,17 @@ 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) - conv_indices, ssm_indices = state_indices - backend: GDNLinearAttBackend = self.backend + backend: LinearAttBackend = 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=conv_indices, + cache_indices=self.b_conv_buffer_idx, has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, @@ -171,7 +190,7 @@ def _gdn_prefill_kernel( # Recurrent processing query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[ssm_indices] + initial_state = ssm_states[self.b_ssm_buffer_idx] # 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, @@ -187,16 +206,90 @@ 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[ssm_indices] = last_recurrent_state.to(ssm_states.dtype, copy=False) + ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) return core_attn_out @dataclasses.dataclass -class GDNLinearAttDecodeAttState(LinearDecodeAttState): - state_class = GDNState +class LinearAttDecodeAttState(BaseDecodeAttState): - def get_draft_step(self): - return self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) + 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): + self._replay_state = None + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import prepare_replay_gdn_decode + + if prepare_replay_gdn_decode(self): + return + 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 decode_att( self, @@ -213,18 +306,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: GDNLinearAttBackend = self.backend + backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - forward = self.get_forward(layer_num) - conv_states, _ = forward.resource.get_layer_view(layer_num) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=type(self._replay_state) if self._replay_state is not None else GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: core_attn_out = self._gdn_mtp_kernel( mixed_qkv, conv_states, - forward, + ssm_states, a, b, self.infer_state, @@ -237,7 +335,7 @@ def decode_att( mixed_qkv, z, conv_states, - forward, + ssm_states, a, b, self.infer_state, @@ -251,15 +349,14 @@ def _gdn_decode_kernel( mixed_qkv: torch.Tensor, z: torch.Tensor, conv_states: torch.Tensor, - forward, + ssm_states: torch.Tensor, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", - *, - layer_num: int, + layer_num: int = 0, ): - backend: GDNLinearAttBackend = self.backend + backend: LinearAttBackend = self.backend # Recurrent processing with fused gating. Decode uses a specialized # conv+pack kernel to avoid materializing the post-conv qkv tensor @@ -272,7 +369,7 @@ def _gdn_decode_kernel( conv_states, layer_weight.linear_conv1d.mm_param.weight, layer_weight.linear_conv1d.bias, - forward.req_indices, + self.b_conv_buffer_idx, backend.activation, backend.conv_kernel_dim, backend.tp_num_k_heads, @@ -280,15 +377,32 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - 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, + if self._replay_state is not None: + core_attn_out = self._replay_state.replay_forward( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + return core_attn_out, z + 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, ) return core_attn_out, z @@ -296,21 +410,20 @@ def _gdn_mtp_kernel( self, mixed_qkv: torch.Tensor, conv_states: torch.Tensor, - forward, + ssm_states: torch.Tensor, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", - *, - layer_num: int, + layer_num: int = 0, ): from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, ) - backend: GDNLinearAttBackend = self.backend + backend: LinearAttBackend = self.backend - cu_seqlens_q = forward.cu_seqlens + cu_seqlens_q = self.b1_mtp_cu_q_seq_len mixed_qkv = causal_conv1d_update_mtp( mixed_qkv, conv_states, @@ -318,20 +431,41 @@ 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=forward.req_indices, - num_accepted_tokens=forward.accepted_token_counts, + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, query_start_loc=cu_seqlens_q, ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - 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, + if self._replay_state is not None: + return self._replay_state.replay_forward( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + cu_seqlens_q, + ) + 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, ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index a985a6bc13..cf443b78bc 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -11,19 +11,24 @@ 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 @@ -44,18 +49,28 @@ def __init__(self, model: "TpPartBaseModel"): self.mtp_step = model.args.mtp_step def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return KDALinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) def create_att_decode_state(self, infer_state: "InferStateInfo"): - return KDALinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDADecodeAttState(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 KDALinearAttPrefillAttState(LinearPrefillAttState): - state_class = KDAState +class KDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + self._replay_state = None + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import prepare_replay_kda_prefill + + if prepare_replay_kda_prefill(self): + return + 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) def prefill_att( self, @@ -74,8 +89,12 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = get_resource(self, layer_num) - conv_indices, ssm_indices = self.state_indices[state.state_name] + 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=type(self._replay_state) if self._replay_state is not None else KDAState, + ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( @@ -83,7 +102,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=conv_indices, + cache_indices=self.b_conv_buffer_idx, has_initial_state=self.infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation="silu", @@ -96,7 +115,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[ssm_indices].contiguous() + initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() output, final_state = chunk_kda_with_fused_gate( q=q, k=k, @@ -112,13 +131,89 @@ def prefill_att( safe_gate=True, lower_bound=backend.lower_bound, ) - ssm_states[ssm_indices] = final_state.to(ssm_states.dtype, copy=False) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) return output @dataclasses.dataclass -class KDALinearAttDecodeAttState(LinearDecodeAttState): - state_class = KDAState +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): + self._replay_state = None + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import prepare_replay_kda_decode + + if prepare_replay_kda_decode(self): + return + 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] def decode_att( self, @@ -137,16 +232,21 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - 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) + 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=type(self._replay_state) if self._replay_state is not None else KDAState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) conv_update = causal_conv1d_update if backend.mtp_step > 0: conv_update = causal_conv1d_update_mtp conv_kwargs.update( mtp_step=backend.mtp_step, - num_accepted_tokens=forward.accepted_token_counts, - query_start_loc=forward.cu_seqlens, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, ) 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) @@ -156,13 +256,32 @@ 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) - 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, + if self._replay_state is not None: + return self._replay_state.replay_forward( + layer_num, + q, + k, + v, + raw_gate.reshape(-1, backend.tp_hidden_size), + raw_beta.reshape(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + 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 output diff --git a/lightllm/common/basemodel/attention/linear/state.py b/lightllm/common/basemodel/attention/linear/state.py deleted file mode 100644 index a3142e6156..0000000000 --- a/lightllm/common/basemodel/attention/linear/state.py +++ /dev/null @@ -1,108 +0,0 @@ -from __future__ import annotations - -import copy -from dataclasses import dataclass, field -from typing import Any - -import torch - -from lightllm.common.attention_state.commit import StateCommit -from lightllm.common.attention_state.bounded_att.ssm.forward import SSMForward -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: Any) -> tuple[Any, ...]: - """返回当前角色下与线性 attention 算法匹配的 bounded 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: Any, logical_layer: int) -> Any: - """按逻辑层号取得当前角色对应的 bounded State。""" - 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) -> None: - """为 prefill 建立每个 State 的请求和物理索引。""" - - 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) - _layout_checked: bool = field(default=False, init=False, repr=False) - - def get_draft_step(self) -> int: - """返回当前模型的 MTP draft 步数。""" - return self.backend.mtp_step - - def init_state(self) -> None: - """根据固定或动态 MTP 布局创建每个 State 的 SSMForward。""" - self._layout_checked = False - 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: int) -> SSMForward: - """返回指定逻辑层共享的 forward 上下文。""" - return self.forwards[get_resource(self, logical_layer).state_name] - - def get_state_commits(self) -> tuple[StateCommit, ...]: - """收集各 State 在本轮 forward 产生的提交凭据。""" - return tuple(commit for forward in self.forwards.values() if (commit := forward.finish()) is not None) - - def clone_for_warmup(self, infer_state: Any) -> LinearDecodeAttState: - """复制 decode 状态供 CUDA Graph warmup 使用。""" - state = super().clone_for_warmup(infer_state) - state.forwards = {name: copy.copy(forward) for name, forward in self.forwards.items()} - state._layout_checked = False - return state - - def copy_for_decode_cuda_graph(self, new_state: LinearDecodeAttState) -> None: - """把动态图输入复制到已 capture 的 CUDA Graph 状态。""" - if not self._layout_checked: - 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]) - self._layout_checked = True diff --git a/lightllm/common/basemodel/attention/linear/triton.py b/lightllm/common/basemodel/attention/linear/triton.py index f579f33a6c..f0b6400f51 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 GDNLinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule -class TritonLinearAttBackend(GDNLinearAttBackend): +class TritonLinearAttBackend(LinearAttBackend): 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 e8e112cc97..0edbdadb94 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -577,18 +577,11 @@ 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 8e4b39688e..d8f2329b24 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -1,14 +1,11 @@ import copy from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional +from typing import 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: @@ -205,9 +202,6 @@ 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 4a046ef02d..ab0e0c0721 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -21,16 +21,6 @@ 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. @@ -133,7 +123,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_warmup_state(infer_state)) + decode_func(copy.copy(infer_state)) torch.cuda.synchronize() for param_name in set(vars(infer_state).keys()): if param_name not in pure_para_set: @@ -166,7 +156,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_warmup_state(infer_state), _copy_warmup_state(infer_state1)) + decode_func(copy.copy(infer_state), copy.copy(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/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index b0f2ea6b79..b81bf9ec16 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 - sequence_mask = offsets < SEQUENCE_CAPACITY - cu_mask = offsets <= SEQUENCE_CAPACITY + cu_mask = offsets <= batch_size 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,24 +25,34 @@ 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) - # A request contributes at most one sequence, regardless of verify width. - # Both the request-pool capacity and graph token bucket are static. + # 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. 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=sequence_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.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]: +) -> tuple[torch.Tensor, 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 @@ -60,41 +70,42 @@ 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. - The fixed-shape outputs are:: + 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:: 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 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. + the padded input batch allows the same CUDA Graph to replay different + per-request query lengths. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda + assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_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((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) + 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) _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 + return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index ae79822bf1..942bca42cd 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -433,7 +433,6 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR state_token_lens=[req.get_chuncked_input_token_len() for req, _ in selected], page_kind="big", page_slot_ids=[slot for _, slot in selected], - flush_history=True, ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -455,7 +454,6 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR state_token_lens=[req.hybrid_cache_len], page_kind="small", page_slot_ids=[dst_buffer_idx], - flush_history=True, ) return 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 5aae49b199..9bc8c16cf7 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,7 +1,6 @@ import torch import time from typing import List -from lightllm.common.attention_state.commit import StateCommitBatch 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,7 +295,7 @@ def decode_mtp( next_token_ids=next_token_ids, b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, - state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), + b_mtp_index=model_input.b_mtp_index, ) 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 1d54f7a21c..ed7e5bb167 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,7 +1,6 @@ import torch import time from typing import List, Tuple -from lightllm.common.attention_state.commit import StateCommitBatch 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,7 +530,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): next_token_ids=next_token_ids, b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, - state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), + b_mtp_index=model_input.b_mtp_index, ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -805,10 +804,7 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf next_token_ids=next_token_ids, b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, - state_commit_batches=( - StateCommitBatch(real_request_num0, model_output0.state_commits), - StateCommitBatch(real_request_num1, model_output1.state_commits), - ), + b_mtp_index=b_mtp_index, ) 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 c32d0a970a..a32cedd17d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,8 +5,6 @@ import torch -from lightllm.common.attention_state.commit import StateAcceptance - from lightllm.common.basemodel.triton_kernel.mtp_utils import ( mtp_scatter_next_token_ids, mtp_verify, @@ -23,9 +21,10 @@ def verify_mtp_tokens( next_token_ids: torch.Tensor, b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, - state_commit_batches=(), + b_mtp_index: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" + # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, @@ -33,8 +32,9 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - acceptance = StateAcceptance(b_req_idx[b_req_mtp_start_loc.long()], accept_lengths, state_commit_batches) - backend.model.att_state_manager.commit_bounded_att_state(acceptance, role="main") + backend.model.att_state_manager.commit_bounded_att_state( + b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" + ) return accept_lengths, accepted_index diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 12bc610eb4..4bf1514704 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -22,7 +22,6 @@ import numpy as np import torch -from lightllm.common.attention_state.commit import StateAcceptance, StateCommitBatch import torch.distributed as dist from transformers import PretrainedConfig @@ -228,24 +227,19 @@ def decode(request: torch.Tensor): max_kv_seq_len=int(lengths[-1]), ) output = model.forward(batch) - commits = [output.state_commits] logits = [output.logits.clone()] for draft in draft_models: batch.input_ids = output.logits.argmax(dim=-1) batch.mtp_draft_input_hiddens = output.mtp_collector.spec_hidden output = draft.forward(batch) - commits.append(output.state_commits) logits.append(output.logits.clone()) if draft_models and hybrid: cuda_req_indices = request.cuda(non_blocking=True) accepted_steps = torch.full_like(cuda_req_indices, accepted_width, dtype=torch.int32) - for item, state_commits in zip((model, *draft_models), commits): + for item in (model, *draft_models): resources.commit_bounded_att_state( - acceptance=StateAcceptance( - cuda_req_indices, - accepted_steps, - (StateCommitBatch(request.numel(), state_commits),), - ), + req_indices=cuda_req_indices, + accepted_steps=accepted_steps, role=item.att_state_role, ) return logits diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 62088a6d91..b531a025b7 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -15,7 +15,19 @@ 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 ReplaySSMExecutor + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.get("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + + return ReplayKDARuntime(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + + return ReplayGDNRuntime(*args, **kwargs) + + from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e7265b3f72..e3a09adb8a 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -14,7 +14,6 @@ import pytest import torch -from lightllm.common.attention_state.commit import StateAcceptance from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3next.model import Qwen3NextTpPartModel @@ -297,10 +296,8 @@ def start_runtime(): draft_states = [manager.get_bounded_attention_states(f"draft{i}")[0] for i in range(draft_count)] assert len({state.conv_state.data_ptr() for state in draft_states}) == draft_count manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), role="main", ) assert all(not state.mtp_state_index.any() for state in draft_states) @@ -309,10 +306,8 @@ def start_runtime(): for i, state in enumerate(draft_states): assert len(state.config.logical_layer_ids) == 1 manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([i + 2], dtype=torch.int32, device="cuda"), - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), role=f"draft{i}", ) assert state.mtp_state_index[0] == i + 1 diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 21c08329e5..ffd59db4aa 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -5,7 +5,6 @@ import pytest import torch -from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ( @@ -255,7 +254,7 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): assert (state.small_pages_pool.conv_state[0] == 7).all() assert (state.small_pages_pool.ssm_state[0] == 11).all() with pytest.raises(AssertionError, match="CUDA tensors"): - state.commit_bounded_att_state(StateAcceptance(torch.tensor([0]), torch.tensor([1], dtype=torch.int32))) + state.commit_bounded_att_state(torch.tensor([0]), torch.tensor([1], dtype=torch.int32)) @pytest.mark.parametrize("cls", [KDAState, GDNState]) diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 7fc7981d32..dc640c029d 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -79,9 +79,9 @@ def test_hybrid_layer_mapping_and_mtp_states( device="meta", ) assert bounded_state.ssm_slots_per_req == (3 if mode == "native" else 1) - assert bounded_state._ssm_executor.state is bounded_state.ssm_state if mode == "replay": - assert bounded_state._ssm_executor.capacity == 8 + assert bounded_state.runtime.state is bounded_state.ssm_state + assert bounded_state.runtime.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 1220f3a5f5..a50ff4fc71 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -7,7 +7,6 @@ import pytest import torch -from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import GDNState @@ -35,9 +34,8 @@ def test_commit_selects_state_for_one_role(): state.ssm_state[:, slot].fill_(slot + 10) state.mtp_state_index.fill_(1) manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="main", ) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.main_gdn_conv_fp32_ssm_fp32_state): @@ -46,9 +44,8 @@ def test_commit_selects_state_for_one_role(): assert (state.ssm_state[:, selected] == 12).all() assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [1, 1, 1] manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([1], dtype=torch.int32, device="cuda") - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([1], dtype=torch.int32, device="cuda"), role="draft0", ) assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 1] @@ -61,22 +58,19 @@ def test_commit_without_recurrent_state_and_invalid_role(): manager = AttentionStateManager(request_capacity=3) manager.finish_state_registration() manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="main", ) manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="draft0", ) with pytest.raises(AssertionError, match="unknown model role"): manager.commit_bounded_att_state( - StateAcceptance( - torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") - ), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="target", ) @@ -86,7 +80,6 @@ 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( @@ -94,7 +87,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,37 +100,10 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i next_token_ids=torch.tensor(tokens, device="cuda"), b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), ) assert accept_lengths.tolist() == [3, 1, 2] assert accepted_index.tolist() == expected_mask assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] - - -def test_verify_routes_one_acceptance(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) - 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 acceptance, **kw: calls.append((acceptance, 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, state_commit_batches=batches) - assert len(calls) == 1 - acceptance, kwargs = calls[0] - assert acceptance.req_indices.tolist() == [2, 0] - assert acceptance.accepted_lengths is lengths - assert acceptance.batches is batches - assert kwargs == {"role": "main"} diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index f9b5b39f76..8ec77c6f4d 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -10,7 +10,6 @@ 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 @@ -39,7 +38,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.GDNLinearAttPrefillAttState( + state = gdn.LinearAttPrefillAttState( backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ) @@ -59,7 +58,6 @@ 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 @@ -71,12 +69,11 @@ 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 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 + 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 @pytest.mark.parametrize("role", ["main", "draft0"]) @@ -92,13 +89,12 @@ def test_decode_state_initializes_fixed_mtp_layout(role: str): ) state.init_state() - forward = state.get_forward(3) - 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.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.state_indices, + state.b_ssm_buffer_idx, torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), ) @@ -114,9 +110,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 + return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens - monkeypatch.setattr(linear_state, "build_dynamic_mtp_linear_att_state_params", build_params) + monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) state, infer_state = _create_decode_state( draft_step=2, dynamic_layout=True, @@ -126,13 +122,12 @@ def build_params(**kwargs): ) state.init_state() - forward = state.get_forward(3) - 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) + 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 torch.testing.assert_close( - forward.state_indices, + state.b_ssm_buffer_idx, torch.tensor( [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], dtype=torch.int32, @@ -141,11 +136,12 @@ 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.GDNLinearAttBackend) + assert inspect.isabstract(gdn.LinearAttBackend) backend = object.__new__(TritonLinearAttBackend) assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule @@ -362,8 +358,6 @@ 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) @@ -385,7 +379,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.GDNLinearAttDecodeAttState( + state = gdn.LinearAttDecodeAttState( 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 3944c0e229..52b1c1cf2d 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -1,6 +1,5 @@ from contextlib import contextmanager from types import SimpleNamespace -import weakref import pytest import torch 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 index cd4927361e..bc51f0fc4a 100644 --- 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 @@ -3,10 +3,23 @@ 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 ReplaySSMExecutor + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.get("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + + return ReplayKDARuntime(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + + return ReplayGDNRuntime(*args, **kwargs) + + from types import SimpleNamespace from lightllm.common.attention_state.commit import StateAcceptance -from lightllm.common.attention_state.bounded_att.ssm.native import NativeSSMExecutor +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix def make(n, width): @@ -22,7 +35,7 @@ def make(n, width): 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 = ReplaySSMExecutor.__new__(ReplaySSMExecutor) + ssm_backend = ReplayGDNRuntime.__new__(ReplayGDNRuntime) ssm_backend.cursors = tensor([7] * (hold + 1)) ssm_backend.hold_req_id = hold ssm_backend.capacity = max(4, triton.next_power_of_2(width)) @@ -117,7 +130,11 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): assert m._ssm_executor.cursors.tolist() == expected_cursor checks += 1 - m._ssm_executor = NativeSSMExecutor(torch.empty(0), width) + class _IndexOnly: + def commit(self, acceptance, indices, state_name=None): + commit_ssm_prefix(indices, acceptance.req_indices, acceptance.accepted_lengths) + + m._ssm_executor = _IndexOnly() accept_fused(m, args) assert m.mtp_state_index.tolist() == expected_index 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 index c4c35a78cf..8252c445a1 100644 --- 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 @@ -29,17 +29,15 @@ def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capa 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 + return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, manager.mtp_state_index, hold) build() torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - cu, ids, ssm, counts = build() - assert cu.numel() == sequences + 1 - assert ids.numel() == counts.numel() == sequences + cu, ids, counts = build() + assert cu.numel() == token_bucket + 1 + assert ids.numel() == counts.numel() == token_bucket # Reuse one graph for full request pools, mixed verify lengths, and HOLD-only padding. for lengths in ([1] * sequences, [min(4, token_bucket // sequences)] * sequences, [2, 1], []): @@ -59,10 +57,9 @@ def build(): 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)) + expected_cu += [tokens] * (token_bucket + 1 - len(expected_cu)) + expected_ids = request_ids + [hold] * (token_bucket - len(request_ids)) assert cu.tolist() == expected_cu assert ids.tolist() == expected_ids - expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (sequences - len(request_ids)) + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (token_bucket - 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 index 568546d602..236d9bff8a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,7 +1,18 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.get("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + + return ReplayKDARuntime(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + + return ReplayGDNRuntime(*args, **kwargs) + + 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") 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 index 42a0899d7f..3f6795b0a6 100644 --- 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 @@ -2,7 +2,18 @@ import torch import triton -from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.get("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + + return ReplayKDARuntime(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + + return ReplayGDNRuntime(*args, **kwargs) + + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) @@ -141,7 +152,7 @@ def inputs(iteration): return ids, lengths, counts, logical def metadata(ssm_backend): - cu, reqs = build_dynamic_mtp_linear_att_state_params(rows, mtp, 4) + cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) return cu, reqs, ssm_backend.prepare_decode(reqs, cu) def forward_accept(index, cu, reqs): 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 index 756774c8b9..4162424509 100644 --- 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 @@ -3,7 +3,17 @@ from lightllm.common.ssm_types import SSMRunConfig from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.get("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + + return ReplayKDARuntime(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + + return ReplayGDNRuntime(*args, **kwargs) @pytest.mark.parametrize("width", [1, 5]) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index cf8cb16c71..2c8163fd2d 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -20,13 +20,18 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): 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_executor is not None: - assert state._ssm_executor.state is state.ssm_state + runtime = getattr(state, "runtime", None) + if runtime is not None: + assert runtime.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) + if hasattr(state, "get_accepted_state"): + actual = state.get_accepted_state(2) + else: + width = state.config.conv_width + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * state.ssm_slots_per_req]) 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 @@ -81,7 +86,7 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): 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_executor is state._ssm_executor + assert getattr(copy, "runtime", None) is getattr(state, "runtime", None) assert copy.mtp_state_index is state.mtp_state_index assert all(vars(state)[name] is value for name, value in before.items()) @@ -105,7 +110,14 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( state._ssm_executor.gates.fill_(-0.5) state._ssm_executor.betas.zero_() state._ssm_executor.cursors[1], state._ssm_executor.cursors[2] = 1, 3 - expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + if hasattr(state, "get_accepted_state"): + expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + else: + width = state.config.conv_width + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, state.ssm_slots_per_req].cpu().tolist(), + ] before = vars(state).copy() source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, state.ssm_slots_per_req].clone() snapshots = [] diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index a5c47924dc..88e9504964 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -19,7 +19,14 @@ def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device= 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}) + if mode == "replay": + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + + cls = ReplayKDAState if cls is KDAState else ReplayGDNState + state = cls(config, 4, mtp_step=step, device=device, ssm_replay_config={"cache_len": 4}) + else: + state = cls(config, 4, mtp_step=step, device=device) resources = AttentionStateManager(request_capacity=4) resources.register_state("main", state) resources.finish_state_registration() @@ -42,8 +49,12 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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 == "native": + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src * state.ssm_slots_per_req].clone() + else: + expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() + expected_ssm = state.ssm_state[:, src].clone() if mode == "replay": state._ssm_executor.raw_keys.zero_() state._ssm_executor.raw_values.zero_() @@ -74,9 +85,14 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if mode == "replay": state._ssm_executor.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 mode == "replay": + actual = state.get_accepted_state(dst) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = state.conv_state[:, dst, ..., :3] + actual_ssm = state.ssm_state[:, dst * state.ssm_slots_per_req] + 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": @@ -111,15 +127,25 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty state._ssm_executor.betas.zero_() state._ssm_executor.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) requests = [2, 0, 1] - expected = [(s.conv.clone(), s.ssm.clone()) for s in map(state.get_accepted_state, requests)] + if mode == "replay": + expected = [(item.conv.clone(), item.ssm.clone()) for item in map(state.get_accepted_state, requests)] + else: + expected = [ + ( + state.conv_state[:, req, :, : state.config.conv_width].clone(), + state.ssm_state[:, req * state.ssm_slots_per_req].clone(), + ) + for req in requests + ] before = state.ssm_state.clone() - cursors = state._ssm_executor.cursors.clone() if mode == "replay" else None + cursors = state.runtime.cursors.clone() if mode == "replay" else None def no_snapshot(*args): raise AssertionError("batch export must write directly to the destination") - monkeypatch.setattr(state, "get_accepted_state", no_snapshot) - monkeypatch.setattr(state._ssm_executor, "get_accepted_state", no_snapshot) + if mode == "replay": + monkeypatch.setattr(state, "get_accepted_state", no_snapshot) + monkeypatch.setattr(state.runtime, "get_accepted_state", no_snapshot) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): @@ -132,13 +158,19 @@ def no_snapshot(*args): assert torch.equal(state.ssm_state, before) if cursors is not None: assert torch.equal(state._ssm_executor.cursors, cursors) - target, _, _ = make_state(cls, "native", 2, dtype) - target.mtp_state_index.fill_(2) + target, _, _ = make_state(cls, mode, 2, dtype) + if target.mtp_state_index is not None: + target.mtp_state_index.fill_(2) target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) for req, (conv, ssm) in enumerate(expected): - actual = target.get_accepted_state(req) - torch.testing.assert_close(actual.conv, conv, atol=0, rtol=0) - torch.testing.assert_close(actual.ssm, ssm, atol=0, rtol=0) + if mode == "replay": + actual = target.get_accepted_state(req) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] + actual_ssm = target.ssm_state[:, req * target.ssm_slots_per_req] + torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) + torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) assert torch.equal(state.ssm_state, before) if cursors is not None: assert torch.equal(state._ssm_executor.cursors, cursors) @@ -171,7 +203,7 @@ def test_native_checkpoint_scratch_is_bounded(runtime_start_args): 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"), + import_module("lightllm.common.attention_state.bounded_att.gdn.gdn_state"), "copy_bounded_att_tensors", copy_on_cpu, ) @@ -193,10 +225,10 @@ def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(mo 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) + state.conv_state[:, src, :2, :3].fill_(10 + rank) + state.conv_state[:, src, 2:4, :3].fill_(20 + rank) + state.conv_state[:, src, 4:, :3].fill_(30 + rank) + state.ssm_state[:, src * state.ssm_slots_per_req].fill_(40 + rank) states.append(state) layouts.append(layout) snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index ca4dc9d3c1..f7223c957e 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -5,10 +5,11 @@ import torch from lightllm.common.ssm_types import SSMRunConfig -from lightllm.common.attention_state.commit import StateAcceptance, 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.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime, ReplayGDNState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -85,14 +86,12 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): def test_kda_backend_receives_gate_bound_at_construction(linear_config): - owner = KDAState(linear_config, 3, mtp_step=2, ssm_state_mode="replay", gate_lower_bound=-2.5, device="cpu") - assert owner._ssm_executor.gate_lower_bound == -2.5 + owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") + assert owner.runtime.gate_lower_bound == -2.5 def test_replay_backend_publishes_detached_run_config(): - from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMExecutor - - ssm_backend = object.__new__(ReplaySSMExecutor) + ssm_backend = object.__new__(ReplayGDNRuntime) tuning = {"BV": 32, "num_warps": 2} ssm_backend.run_config = tuning ssm_backend._config_is_fixed = True @@ -109,134 +108,36 @@ def test_replay_backend_publishes_detached_run_config(): @pytest.mark.parametrize("cls", [GDNState, KDAState]) @pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): - owner = cls(linear_config, request_capacity=2, mtp_step=mtp_step, ssm_state_mode="replay", device="cpu") - assert owner._ssm_executor.compute_mode == expected - - -def make_acceptance(batches=()): - return StateAcceptance(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), batches) - - -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)) - acceptance = make_acceptance((StateCommitBatch(3, (commit,)),)) - with pytest.raises(ValueError, match="unknown state"): - acceptance.validate({"draft0"}) - with pytest.raises(ValueError, match="layout"): - StateAcceptance(torch.tensor([1, 2]), torch.tensor([1])).validate() - with pytest.raises(ValueError, match="partition"): - make_acceptance((StateCommitBatch(2),)).validate({"main"}) + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") + assert owner.runtime.compute_mode == expected @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(StateAcceptance(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() +def test_native_commit_selects_mtp_slot(cls): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(cls, step=2) + owner.commit_bounded_att_state( + torch.tensor([1], device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ) 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 + import sys owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") calls = [] - monkeypatch.setattr(base, "commit_ssm_prefix", lambda *args: calls.append(args)) - monkeypatch.setattr(owner._ssm_executor, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) - owner.commit_bounded_att_state(make_acceptance()) - assert len(calls) == 1 - assert calls[0][3] is owner._ssm_executor.cursors - assert calls[0][0] is owner.mtp_state_index - - -@pytest.mark.parametrize("change", ["owner", "executor", "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 == "executor": - source.executor = copy(owner._ssm_executor) - 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] - - -@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 + sys.modules[owner.runtime.__class__.__module__], "commit_ssm_prefix", lambda *args: calls.append(args) ) - 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 + monkeypatch.setattr(owner.runtime, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3])) + assert len(calls) == 1 + assert calls[0][3] is owner.runtime.cursors + assert calls[0][0] is owner.mtp_state_index @pytest.mark.parametrize("device", ["cpu", "cuda"]) @@ -245,9 +146,9 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): pytest.skip("CUDA required") owner, _, _ = make_state(mode="replay", step=3, device=device) owner.mtp_state_index.fill_(2) - owner._ssm_executor.cursors.fill_(8) + owner.runtime.cursors.fill_(8) reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] - owner.commit_bounded_att_state(StateAcceptance(reqs, lengths)) + owner.commit_bounded_att_state(reqs, lengths) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] - assert owner._ssm_executor.cursors.tolist() == [11, 8, 9, 8] + assert owner.runtime.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index fdcb28beab..f1a7a151f9 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -130,7 +130,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update torch.manual_seed(53) @@ -206,7 +207,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod ) selected = None for model_role in ("main", "draft0"): - linear = KDAState(config, request_capacity=4, mtp_step=mtp_step, device="cuda", ssm_state_mode=mode) + state_cls = ReplayKDAState if mode == "replay" else KDAState + linear = state_cls(config, request_capacity=4, mtp_step=mtp_step, device="cuda") if mtp_step: linear.mtp_state_index.copy_(accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets)) if model_role == role: @@ -232,7 +234,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDALinearAttDecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() actual = state.decode_att( None, diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index a2af828623..e020c6709b 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -290,12 +290,7 @@ 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, - decode_att_state=None, - decode_att_state1=None, - ) + state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=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]) From ed4f14086c75aa030ca7c850f1a7a5a85ad8d7f0 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 02:04:09 +0000 Subject: [PATCH 07/54] refactor: move GDN and KDA replay off the native attention path Native linear backends only run fused recurrent now. Replay keeps its own state modules and attention backends, selected when ssm_state_mode is replay. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_checkpoint.py | 296 +++++ .../bounded_att/gdn/replay_gdn_runtime.py | 425 +++++++ .../bounded_att/gdn/replay_gdn_snapshot.py | 117 ++ .../bounded_att/gdn/replay_gdn_state.py | 1059 ++--------------- .../bounded_att/kda/replay_kda_checkpoint.py | 296 +++++ .../bounded_att/kda/replay_kda_runtime.py | 425 +++++++ .../bounded_att/kda/replay_kda_snapshot.py | 117 ++ .../bounded_att/kda/replay_kda_state.py | 1059 ++--------------- .../attention/create_linear_utils.py | 26 +- .../common/basemodel/attention/linear/gdn.py | 45 +- .../common/basemodel/attention/linear/kda.py | 27 +- .../basemodel/attention/linear/replay_gdn.py | 268 +++++ .../basemodel/attention/linear/replay_kda.py | 226 ++++ lightllm/models/glm5_next/model.py | 19 +- test/benchmark/kernels/benchmark_kda_ssm.py | 4 +- .../linear_att/test_acceptance_fusion.py | 6 +- .../linear_att/test_replayssm.py | 4 +- .../linear_att/test_replayssm_execution.py | 4 +- .../linear_att/test_ssm_autotune_memory.py | 4 +- .../common/test_linear_state_contract.py | 3 +- unit_tests/models/glm5_next/test_mtp.py | 55 +- 21 files changed, 2504 insertions(+), 1981 deletions(-) create mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py create mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py create mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py create mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py create mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py create mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_kda.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py new file mode 100644 index 0000000000..99f197b40b --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py @@ -0,0 +1,296 @@ +"""GDN replay 的页布局和快照拷贝。 + +这里只处理 Conv/SSM 快照页，不读取 replay cursor 或 history。 +""" + +from __future__ import annotations + +from math import prod +from typing import List + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...types import ( + AxisMapping, + PageKind, + StateTransferLayout, + TensorField, + TransferRegion, +) +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .gdn_pages_pool import GDNPagesPool + +logger = init_logger(__name__) + + +def build_transfer_regions(state) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = state.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", + state.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", + state.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)", + state.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + +def page_pool(state, page_kind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {state.state_name}" + return pool + + +def tensor_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")), + } + + +def write_pages_to_cpu( + state, + 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 == state.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 = page_pool(state, 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(tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + +def read_pages_from_cpu( + state, + 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 == state.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 = state.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 = [tensor_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( + state, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state, +) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = state.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 init_page_pools( + state, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, +) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert state.big_pages_pool is None and state.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + state.big_pages_pool = ( + GDNPagesPool( + size=big_page_capacity, + config=state.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + state.small_pages_pool = ( + GDNPagesPool( + size=small_page_capacity, + config=state.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + +# ======================================================================== +# 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 +# ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py new file mode 100644 index 0000000000..828d535c14 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py @@ -0,0 +1,425 @@ +"""GDN replay history。只负责 ring buffer、重放 kernel 和 accepted SSM 物化。""" + +from __future__ import annotations + +from collections.abc import Sequence + +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 lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( + commit_ssm_prefix, +) + + +def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +class ReplayGDNRuntime: + """以短 history 重放 decode 的 SSM 执行器。 + + ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 + cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode + 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 + checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 + """ + + _run_config: SSMRunConfig | None = None + + @property + def run_config(self) -> SSMRunConfig | None: + """返回最近一次 forward 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) + + def __init__( + self, + state: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = False, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> None: + """创建 GDN replay 的 canonical state、双半区 history 和 cursor。 + + ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 + """ + + assert state.dtype in ( + torch.float32, + torch.bfloat16, + ), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + if kda: + raise ValueError("GDN replay runtime does not use the KDA gate") + self.state = state + self.kda = False + 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 + # 保留冗余布局，兼容已有调用方和 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) + # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 + # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 + # 下可以使用更窄的 history buffer；其他形状使用 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), + 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: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + 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, + 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, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor 以保留每个请求本轮已接受的前缀。 + + MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft + 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 + """ + _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: torch.Tensor) -> None: + """把接受的 history 物化回 canonical state，并清空请求 cursor。""" + 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: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + + layers, slots, hv, k, v = self.state.shape + config = self.run_config or SSMRunConfig(32, 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, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: + """返回不修改 active state/history 的单请求 canonical 快照。""" + 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 export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" + + if not req_indices: + return + shape = (self.state.shape[0], *self.state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) + self._materialize_accepted_state( + upload_indices(req_indices, self.state.device, torch.int32), + upload_indices([t.data_ptr() for t in targets], self.state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 GDN replay 调优配置。""" + + return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward( + self, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, + *, + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 GDN replay decode，并返回输出与本轮 run config。""" + + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + 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 == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor + # 运算，无法定义静态 shape。 + 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, + self.kda, + self.gate_lower_bound, + 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: torch.Tensor) -> None: + """非 MTP decode 直接把单 token 写入当前 history。""" + self.accept_updates(reqs) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm( + self, + req_indices: list[int], + targets: list[torch.Tensor], + *, + flush: bool = False, + ) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.state.shape[0], *self.state.shape[2:]), + dtype=self.state.dtype, + device=self.state.device, + ) + self._materialize_accepted_state( + torch.tensor([req], dtype=torch.int32, device=self.state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.state[:, req].copy_(snap) + self.clear_history(req) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py new file mode 100644 index 0000000000..9eaecd0d2e --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py @@ -0,0 +1,117 @@ +"""GDN replay 的已接受快照读写。 + +页布局在 replay_gdn_checkpoint。这里决定写哪一段 Conv、是否物化 SSM，以及何时清 history。 +""" + +from __future__ import annotations + +import torch + +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .replay_gdn_checkpoint import page_pool, tensor_fields + + +def accepted_conv(state, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = state.config.conv_width + conv = state.conv_state[:, req_idx] + if state.mtp_state_index is None: + return conv[:, :, :width] + offset = state.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + +def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = page_pool(state, page_kind) + copy_bounded_att_tensors( + copies=[ + (accepted_conv(state, req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + state.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + +def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = page_pool(state, page_kind) + width = state.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + state.runtime.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if state.mtp_state_index is not None: + for req_idx in req_indices: + state.mtp_state_index[req_idx].fill_(0) + + +def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_bounded_state) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = state.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + state.runtime.clear_history(dst_req) + conv_copies.append( + ( + accepted_conv(remote_dp_bounded_state, src_req), + state.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.runtime.export_ssm( + src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices] + ) + if state.mtp_state_index is not None: + for dst_req in dst_req_indices: + state.mtp_state_index[dst_req].fill_(0) + + +def write_pd(state, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == state.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = state.get_accepted_state(req_idx) + fields = tensor_fields(accepted.conv, accepted.ssm) + 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_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == state.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + state.runtime.clear_history(req_idx) + fields = tensor_fields( + state.conv_state[:, req_idx, :, : state.config.conv_width], + state.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if state.mtp_state_index is not None: + state.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 5954b3846f..4c1c9ac2e7 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -1,443 +1,51 @@ -"""GDN replay 运行态。 +"""GDN replay 的 bounded attention 状态。 -这个文件单独实现 GDN 的 replay history、checkpoint 和 decode 准备。 -它不继承 native ``GDNState``，也不和 KDA replay 共享执行器基类。 +直接实现状态接口，不继承 native GDNState。history 在 replay_gdn_runtime， +页布局在 replay_gdn_checkpoint，已接受快照在 replay_gdn_snapshot。 """ from __future__ import annotations -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 collections.abc import Sequence - -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix - - -def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -class ReplayGDNRuntime: - """以短 history 重放 decode 的 SSM 执行器。 - - ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 - cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode - 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 - checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 - """ - - _run_config: SSMRunConfig | None = None - - @property - def run_config(self) -> SSMRunConfig | None: - """返回最近一次 forward 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" - - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) - - def __init__( - self, - state: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - *, - num_key_heads: int | None = None, - kda: bool = False, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, - ) -> None: - """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 - - ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 - """ - - 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") - 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 - # 保留冗余布局，兼容已有调用方和 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) - # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 - # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 - # 下可以使用更窄的 history buffer；其他形状使用 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: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - 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, - 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, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor 以保留每个请求本轮已接受的前缀。 - - MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft - 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 - """ - _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: torch.Tensor) -> None: - """把接受的 history 物化回 canonical state，并清空请求 cursor。""" - 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: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - - layers, slots, hv, k, v = self.state.shape - config = self.run_config or SSMRunConfig(32, 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, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回不修改 active state/history 的单请求 canonical 快照。""" - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" - - if not req_indices: - return - shape = (self.state.shape[0], *self.state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) - self._materialize_accepted_state( - upload_indices(req_indices, self.state.device, torch.int32), - upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> SSMRunConfig: - """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" - - return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig]: - """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" - - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - 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 == "precompute" - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor - # 运算，无法定义静态 shape。 - 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, - self.kda, - self.gate_lower_bound, - 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: torch.Tensor) -> None: - """非 MTP decode 直接把单 token 写入当前 history。""" - self.accept_updates(reqs) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) - - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" - - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.state.shape[0], *self.state.shape[2:]), - dtype=self.state.dtype, - device=self.state.device, - ) - self._materialize_accepted_state( - torch.tensor([req], dtype=torch.int32, device=self.state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.state[:, req].copy_(snap) - self.clear_history(req) - - from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod +from typing import NamedTuple import torch -from lightllm.utils.log_utils import init_logger - from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig -from .gdn_pages_pool import GDNPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind +from .replay_gdn_checkpoint import ( + build_transfer_regions, + copy_big_pages, + init_page_pools, + read_pages_from_cpu, + write_pages_to_cpu, +) +from .replay_gdn_runtime import ReplayGDNRuntime +from .replay_gdn_snapshot import ( + accepted_conv, + copy_runtime_from_dp, + read_pd, + restore_runtime, + save_runtime, + write_pd, +) -logger = init_logger(__name__) +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor -class ReplayGDNState(BoundedAttentionState): - """GDN replay：每个请求一份 canonical SSM，外加一段不进入 checkpoint 的 history。""" +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode - big_pages_pool: GDNPagesPool | None - small_pages_pool: GDNPagesPool | None - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== +class ReplayGDNState(BoundedAttentionState): + """GDN replay 运行态。槽位布局和 history 都在本类，不沿用 native 的 MTP 槽。""" def __init__( self, @@ -450,28 +58,31 @@ def __init__( ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, ) -> None: - """按请求容量创建 Conv 运行态、单槽 canonical SSM 和 replay history。""" + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" if ssm_state_mode != "replay": raise ValueError(f"ReplayGDNState only implements replay mode, got {ssm_state_mode}") 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.state_name = None + self.request_capacity = request_capacity self.mtp_step = mtp_step self.gate_lower_bound = gate_lower_bound - 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._algorithm = "gdn" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + ( + 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, + layer_count, request_capacity, config.num_v_heads, config.key_dim, @@ -480,6 +91,9 @@ def __init__( dtype=config.ssm_dtype, device=device, ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) self.runtime = ReplayGDNRuntime( self.ssm_state, @@ -491,25 +105,25 @@ def __init__( gate_lower_bound=gate_lower_bound, compute_mode=compute_mode, ) - 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 @property def ssm_slots_per_req(self) -> int: - """Replay 每个请求只保留一个 canonical SSM 槽位。""" + """每个请求只保留一份 canonical SSM。""" return 1 @property def _ssm_executor(self): - """供既有测试读取 history；实现本身不依赖这个名字。""" + """测试沿用的 history 入口。""" return self.runtime 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)}" + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) def initialize_page_pools( self, @@ -518,45 +132,52 @@ def initialize_page_pools( 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 提交 - # ======================================================================== + init_page_pools(self, big_page_capacity, small_page_capacity, device, pin_memory) 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 get_layer_view(self, logical_layer_id: int): + index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[index], self.ssm_state[index] - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + return build_transfer_regions(self) + + def write_state_pages_to_cpu( + self, + page_kind, + page_slot_ids, + cpu_page_buffers, + cpu_page_indices, + layout, + cpu_page_ready_flags=None, + ): + write_pages_to_cpu( + self, + page_kind, + page_slot_ids, + cpu_page_buffers, + cpu_page_indices, + layout, + cpu_page_ready_flags, + ) + + def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): + read_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout) + + def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): + copy_big_pages(self, src_page_ids, dst_page_ids, remote_dp_bounded_state) + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self): + return build_transfer_regions(self) + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" assert isinstance(req_indices, list), "request indices must be a list" for req_idx in req_indices: self.conv_state[:, req_idx].zero_() @@ -567,67 +188,34 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], + req_indices: list[int], + state_token_lens: list[int], + page_kind, + 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 = [ - (self._accepted_conv(req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - copy_bounded_att_tensors(copies=conv_copies) - self.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + save_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], + req_indices: list[int], + state_token_lens: list[int], + page_kind, + 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) - width = self.config.conv_width - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - self.runtime.clear_history(req_idx) - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) + restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """提交已接受前缀：更新 MTP 游标，并把同样的长度加到 replay cursor。""" + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受前缀，并推进 replay cursor。""" if self.mtp_state_index is None: return self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """Prefill 前把已接受 history 折回 canonical state。""" - self.runtime.merge_accepted_updates(reqs) def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。非 MTP decode 同时接受本轮唯一 token。""" - + """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" self.runtime.prepare_decode(reqs, cu_seqlens) if cu_seqlens is None: self.runtime.accept_single_token(reqs) @@ -645,8 +233,7 @@ def replay_forward( reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: - """执行一层 GDN replay decode。""" - + """执行一层 replay decode。""" output, _ = self.runtime.forward( self._logical_to_storage_layer[logical_layer], q, @@ -661,165 +248,17 @@ def replay_forward( ) return output - def get_accepted_state(self, req_idx: int): - """返回已接受 Conv 窗口和物化后的 SSM，不改写 canonical state 或 cursor。""" - - ssm = self.runtime.get_accepted_state(req_idx) - return _AcceptedTensors(self._accepted_conv(req_idx), ssm) - - def _accepted_conv(self, req_idx: int) -> torch.Tensor: - width = self.config.conv_width - conv = self.conv_state[:, req_idx] - if self.mtp_state_index is None: - return conv[:, :, :width] - offset = self.mtp_state_index[req_idx] - index = torch.arange(width, device=conv.device) + offset - return conv.index_select(-1, index) - - # ======================================================================== - # 三、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: ReplayGDNState, - ) -> 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 get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) def copy_runtime_state_from_dp( self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: ReplayGDNState, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayGDNState", ) -> 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" - width = self.config.conv_width - conv_copies = [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.runtime.clear_history(dst_req) - conv_copies.append( - ( - remote_dp_bounded_state._accepted_conv(src_req), - self.conv_state[:, dst_req, :, :width], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.runtime.export_ssm( - src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices] - ) - 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() + copy_runtime_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state) def write_to_pd_page( self, @@ -827,21 +266,9 @@ def write_to_pd_page( pd_page_buffers: torch.Tensor, pd_page_index: int, state_token_len: int, - layout: StateTransferLayout, + layout, ) -> None: - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = self.get_accepted_state(req_idx) - fields = self._fields(accepted.conv, accepted.ssm) - 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) + write_pd(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout) def read_from_pd_page( self, @@ -849,254 +276,6 @@ def read_from_pd_page( req_idx: int, pd_page_index: int, state_token_len: int, - layout: StateTransferLayout, + layout, ) -> 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" - self.runtime.clear_history(req_idx) - fields = self._fields( - self.conv_state[:, req_idx, :, : self.config.conv_width], - self.ssm_state[:, req_idx], - ) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - 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]) - - @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")), - } - - -from typing import NamedTuple - - -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode - - -def _find_replay_state(att_state): - manager = getattr(att_state.infer_state, "att_state_manager", None) - if manager is None: - return None - role = att_state.infer_state.att_state_role - matches = [state for state in manager.get_bounded_attention_states(role) if isinstance(state, ReplayGDNState)] - if not matches: - return None - if len(matches) != 1: - raise ValueError(f"expected one ReplayGDNState for {role}, found {len(matches)}") - return matches[0] - - -def prepare_replay_gdn_prefill(att_state) -> bool: - """Replay prefill 先折叠 history，并用单槽 SSM 索引。""" - - state = _find_replay_state(att_state) - if state is None: - return False - reqs = att_state.infer_state.b_req_idx - state.merge_accepted_updates(reqs) - att_state.b_conv_buffer_idx = reqs - att_state.b_ssm_buffer_idx = reqs - att_state._replay_state = state - return True - - -def prepare_replay_gdn_decode(att_state) -> bool: - """准备 GDN replay decode 的 conv/MTP 索引，并折叠本轮 history。""" - - state = _find_replay_state(att_state) - if state is None: - return False - draft_step = att_state.backend.model.mtp_manager.get_decode_draft_step(att_state.backend.model.is_mtp_draft_model) - infer = att_state.infer_state - if draft_step == 0: - att_state.b_conv_buffer_idx = infer.b_req_idx - att_state.b_ssm_buffer_idx = infer.b_req_idx - att_state.b1_mtp_cu_q_seq_len = None - att_state.b_num_accepted_tokens = None - else: - if att_state.backend.uses_dynamic_spec_verify_layout(): - from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, - ) - - ( - att_state.b1_mtp_cu_q_seq_len, - att_state.b_conv_buffer_idx, - att_state.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=state.mtp_state_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - else: - width = draft_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "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_state.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - att_state.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - att_state.b_num_accepted_tokens = state.mtp_state_index[att_state.b_conv_buffer_idx] + 1 - att_state.b_ssm_buffer_idx = att_state.b_conv_buffer_idx - state.prepare_for_decode(att_state.b_conv_buffer_idx, att_state.b1_mtp_cu_q_seq_len) - att_state._replay_state = state - return True + read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py new file mode 100644 index 0000000000..5723b71db9 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py @@ -0,0 +1,296 @@ +"""KDA replay 的页布局和快照拷贝。 + +这里只处理 Conv/SSM 快照页，不读取 replay cursor 或 history。 +""" + +from __future__ import annotations + +from math import prod +from typing import List + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...types import ( + AxisMapping, + PageKind, + StateTransferLayout, + TensorField, + TransferRegion, +) +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .kda_pages_pool import KDAPagesPool + +logger = init_logger(__name__) + + +def build_transfer_regions(state) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = state.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", + state.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", + state.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)", + state.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + +def page_pool(state, page_kind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {state.state_name}" + return pool + + +def tensor_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")), + } + + +def write_pages_to_cpu( + state, + 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 == state.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 = page_pool(state, 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(tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + +def read_pages_from_cpu( + state, + 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 == state.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 = state.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 = [tensor_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( + state, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state, +) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = state.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 init_page_pools( + state, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, +) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert state.big_pages_pool is None and state.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + state.big_pages_pool = ( + KDAPagesPool( + size=big_page_capacity, + config=state.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + state.small_pages_pool = ( + KDAPagesPool( + size=small_page_capacity, + config=state.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + +# ======================================================================== +# 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 +# ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py new file mode 100644 index 0000000000..53c2f75447 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py @@ -0,0 +1,425 @@ +"""KDA replay history。只负责 ring buffer、重放 kernel 和 accepted SSM 物化。""" + +from __future__ import annotations + +from collections.abc import Sequence + +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 lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( + commit_ssm_prefix, +) + + +def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +class ReplayKDARuntime: + """以短 history 重放 decode 的 SSM 执行器。 + + ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 + cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode + 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 + checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 + """ + + _run_config: SSMRunConfig | None = None + + @property + def run_config(self) -> SSMRunConfig | None: + """返回最近一次 forward 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) + + def __init__( + self, + state: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = True, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> None: + """创建 KDA replay 的 canonical state、双半区 history 和 cursor。 + + ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 + history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 + """ + + assert state.dtype in ( + torch.float32, + torch.bfloat16, + ), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + if not kda: + raise ValueError("KDA replay runtime requires the KDA gate") + self.state = state + self.kda = True + 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 + # 保留冗余布局，兼容已有调用方和 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) + # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 + # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 + # 下可以使用更窄的 history buffer；其他形状使用 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), + 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: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + 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, + 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, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor 以保留每个请求本轮已接受的前缀。 + + MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft + 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 + """ + _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: torch.Tensor) -> None: + """把接受的 history 物化回 canonical state，并清空请求 cursor。""" + 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: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + + layers, slots, hv, k, v = self.state.shape + config = self.run_config or SSMRunConfig(32, 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, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: + """返回不修改 active state/history 的单请求 canonical 快照。""" + 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 export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" + + if not req_indices: + return + shape = (self.state.shape[0], *self.state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) + self._materialize_accepted_state( + upload_indices(req_indices, self.state.device, torch.int32), + upload_indices([t.data_ptr() for t in targets], self.state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + ) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 KDA replay 调优配置。""" + + return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward( + self, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, + *, + state_indices: torch.Tensor | None = None, + accepted_token_counts: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 KDA replay decode，并返回输出与本轮 run config。""" + + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + 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 == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor + # 运算，无法定义静态 shape。 + 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, + self.kda, + self.gate_lower_bound, + 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: torch.Tensor) -> None: + """非 MTP decode 直接把单 token 写入当前 history。""" + self.accept_updates(reqs) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm( + self, + req_indices: list[int], + targets: list[torch.Tensor], + *, + flush: bool = False, + ) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.state.shape[0], *self.state.shape[2:]), + dtype=self.state.dtype, + device=self.state.device, + ) + self._materialize_accepted_state( + torch.tensor([req], dtype=torch.int32, device=self.state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.state[:, req].copy_(snap) + self.clear_history(req) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py new file mode 100644 index 0000000000..d57dc65038 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py @@ -0,0 +1,117 @@ +"""KDA replay 的已接受快照读写。 + +页布局在 replay_kda_checkpoint。这里决定写哪一段 Conv、是否物化 SSM，以及何时清 history。 +""" + +from __future__ import annotations + +import torch + +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .replay_kda_checkpoint import page_pool, tensor_fields + + +def accepted_conv(state, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = state.config.conv_width + conv = state.conv_state[:, req_idx] + if state.mtp_state_index is None: + return conv[:, :, :width] + offset = state.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + +def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = page_pool(state, page_kind) + copy_bounded_att_tensors( + copies=[ + (accepted_conv(state, req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + state.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + +def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = page_pool(state, page_kind) + width = state.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + state.runtime.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if state.mtp_state_index is not None: + for req_idx in req_indices: + state.mtp_state_index[req_idx].fill_(0) + + +def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_bounded_state) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = state.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + state.runtime.clear_history(dst_req) + conv_copies.append( + ( + accepted_conv(remote_dp_bounded_state, src_req), + state.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.runtime.export_ssm( + src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices] + ) + if state.mtp_state_index is not None: + for dst_req in dst_req_indices: + state.mtp_state_index[dst_req].fill_(0) + + +def write_pd(state, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == state.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = state.get_accepted_state(req_idx) + fields = tensor_fields(accepted.conv, accepted.ssm) + 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_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == state.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + state.runtime.clear_history(req_idx) + fields = tensor_fields( + state.conv_state[:, req_idx, :, : state.config.conv_width], + state.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if state.mtp_state_index is not None: + state.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 2b4e33697c..0e877746a8 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -1,443 +1,51 @@ -"""KDA replay 运行态。 +"""KDA replay 的 bounded attention 状态。 -这个文件单独实现 KDA 的 replay history、checkpoint 和 decode 准备。 -它不继承 native ``GDNState``，也不和 GDN replay 共享执行器基类。 +直接实现状态接口，不继承 native KDAState。history 在 replay_kda_runtime， +页布局在 replay_kda_checkpoint，已接受快照在 replay_kda_snapshot。 """ from __future__ import annotations -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 collections.abc import Sequence - -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix - - -def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -class ReplayKDARuntime: - """以短 history 重放 decode 的 SSM 执行器。 - - ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 - cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode - 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 - checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 - """ - - _run_config: SSMRunConfig | None = None - - @property - def run_config(self) -> SSMRunConfig | None: - """返回最近一次 forward 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" - - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) - - def __init__( - self, - state: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - *, - num_key_heads: int | None = None, - kda: bool = True, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, - ) -> None: - """创建 ReplaySSM 的 canonical state、双半区 history 和 cursor。 - - ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 - """ - - 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") - 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 - # 保留冗余布局，兼容已有调用方和 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) - # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 - # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 - # 下可以使用更窄的 history buffer；其他形状使用 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: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - 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, - 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, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor 以保留每个请求本轮已接受的前缀。 - - MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft - 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 - """ - _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: torch.Tensor) -> None: - """把接受的 history 物化回 canonical state，并清空请求 cursor。""" - 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: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - - layers, slots, hv, k, v = self.state.shape - config = self.run_config or SSMRunConfig(32, 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, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回不修改 active state/history 的单请求 canonical 快照。""" - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" - - if not req_indices: - return - shape = (self.state.shape[0], *self.state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) - self._materialize_accepted_state( - upload_indices(req_indices, self.state.device, torch.int32), - upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> SSMRunConfig: - """按当前 kernel 形状选择或返回固定的 ReplaySSM 调优配置。""" - - return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig]: - """执行一层 ReplaySSM decode，并返回输出与本轮 run config。""" - - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - 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 == "precompute" - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor - # 运算，无法定义静态 shape。 - 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, - self.kda, - self.gate_lower_bound, - 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: torch.Tensor) -> None: - """非 MTP decode 直接把单 token 写入当前 history。""" - self.accept_updates(reqs) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) - - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" - - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.state.shape[0], *self.state.shape[2:]), - dtype=self.state.dtype, - device=self.state.device, - ) - self._materialize_accepted_state( - torch.tensor([req], dtype=torch.int32, device=self.state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.state[:, req].copy_(snap) - self.clear_history(req) - - from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod +from typing import NamedTuple import torch -from lightllm.utils.log_utils import init_logger - from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig -from .kda_pages_pool import KDAPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind +from .replay_kda_checkpoint import ( + build_transfer_regions, + copy_big_pages, + init_page_pools, + read_pages_from_cpu, + write_pages_to_cpu, +) +from .replay_kda_runtime import ReplayKDARuntime +from .replay_kda_snapshot import ( + accepted_conv, + copy_runtime_from_dp, + read_pd, + restore_runtime, + save_runtime, + write_pd, +) -logger = init_logger(__name__) +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor -class ReplayKDAState(BoundedAttentionState): - """KDA replay：每个请求一份 canonical SSM，外加一段不进入 checkpoint 的 history。""" +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode - big_pages_pool: KDAPagesPool | None - small_pages_pool: KDAPagesPool | None - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== +class ReplayKDAState(BoundedAttentionState): + """KDA replay 运行态。槽位布局和 history 都在本类，不沿用 native 的 MTP 槽。""" def __init__( self, @@ -450,28 +58,31 @@ def __init__( ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, ) -> None: - """按请求容量创建 Conv 运行态、单槽 canonical SSM 和 replay history。""" + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" if ssm_state_mode != "replay": raise ValueError(f"ReplayKDAState only implements replay mode, got {ssm_state_mode}") 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.state_name = None + self.request_capacity = request_capacity self.mtp_step = mtp_step self.gate_lower_bound = gate_lower_bound - 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._algorithm = "kda" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + ( + 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, + layer_count, request_capacity, config.num_v_heads, config.key_dim, @@ -480,6 +91,9 @@ def __init__( dtype=config.ssm_dtype, device=device, ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) self.runtime = ReplayKDARuntime( self.ssm_state, @@ -491,25 +105,25 @@ def __init__( gate_lower_bound=gate_lower_bound, compute_mode=compute_mode, ) - 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 @property def ssm_slots_per_req(self) -> int: - """Replay 每个请求只保留一个 canonical SSM 槽位。""" + """每个请求只保留一份 canonical SSM。""" return 1 @property def _ssm_executor(self): - """供既有测试读取 history；实现本身不依赖这个名字。""" + """测试沿用的 history 入口。""" return self.runtime def get_state_type(self) -> str: - """返回 GDN 算法及 Conv/SSM 存储精度。""" - return f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) def initialize_page_pools( self, @@ -518,45 +132,52 @@ def initialize_page_pools( 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 提交 - # ======================================================================== + init_page_pools(self, big_page_capacity, small_page_capacity, device, pin_memory) 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 get_layer_view(self, logical_layer_id: int): + index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[index], self.ssm_state[index] - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + return build_transfer_regions(self) + + def write_state_pages_to_cpu( + self, + page_kind, + page_slot_ids, + cpu_page_buffers, + cpu_page_indices, + layout, + cpu_page_ready_flags=None, + ): + write_pages_to_cpu( + self, + page_kind, + page_slot_ids, + cpu_page_buffers, + cpu_page_indices, + layout, + cpu_page_ready_flags, + ) + + def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): + read_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout) + + def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): + copy_big_pages(self, src_page_ids, dst_page_ids, remote_dp_bounded_state) + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self): + return build_transfer_regions(self) + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" assert isinstance(req_indices, list), "request indices must be a list" for req_idx in req_indices: self.conv_state[:, req_idx].zero_() @@ -567,67 +188,34 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], + req_indices: list[int], + state_token_lens: list[int], + page_kind, + 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 = [ - (self._accepted_conv(req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - copy_bounded_att_tensors(copies=conv_copies) - self.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + save_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], + req_indices: list[int], + state_token_lens: list[int], + page_kind, + 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) - width = self.config.conv_width - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - self.runtime.clear_history(req_idx) - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) + restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """提交已接受前缀：更新 MTP 游标，并把同样的长度加到 replay cursor。""" + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受前缀，并推进 replay cursor。""" if self.mtp_state_index is None: return self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """Prefill 前把已接受 history 折回 canonical state。""" - self.runtime.merge_accepted_updates(reqs) def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。非 MTP decode 同时接受本轮唯一 token。""" - + """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" self.runtime.prepare_decode(reqs, cu_seqlens) if cu_seqlens is None: self.runtime.accept_single_token(reqs) @@ -645,8 +233,7 @@ def replay_forward( reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: - """执行一层 KDA replay decode。""" - + """执行一层 replay decode。""" output, _ = self.runtime.forward( self._logical_to_storage_layer[logical_layer], q, @@ -661,165 +248,17 @@ def replay_forward( ) return output - def get_accepted_state(self, req_idx: int): - """返回已接受 Conv 窗口和物化后的 SSM，不改写 canonical state 或 cursor。""" - - ssm = self.runtime.get_accepted_state(req_idx) - return _AcceptedTensors(self._accepted_conv(req_idx), ssm) - - def _accepted_conv(self, req_idx: int) -> torch.Tensor: - width = self.config.conv_width - conv = self.conv_state[:, req_idx] - if self.mtp_state_index is None: - return conv[:, :, :width] - offset = self.mtp_state_index[req_idx] - index = torch.arange(width, device=conv.device) + offset - return conv.index_select(-1, index) - - # ======================================================================== - # 三、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: ReplayKDAState, - ) -> 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 get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) def copy_runtime_state_from_dp( self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: ReplayKDAState, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayKDAState", ) -> 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" - width = self.config.conv_width - conv_copies = [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.runtime.clear_history(dst_req) - conv_copies.append( - ( - remote_dp_bounded_state._accepted_conv(src_req), - self.conv_state[:, dst_req, :, :width], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.runtime.export_ssm( - src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices] - ) - 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() + copy_runtime_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state) def write_to_pd_page( self, @@ -827,21 +266,9 @@ def write_to_pd_page( pd_page_buffers: torch.Tensor, pd_page_index: int, state_token_len: int, - layout: StateTransferLayout, + layout, ) -> None: - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = self.get_accepted_state(req_idx) - fields = self._fields(accepted.conv, accepted.ssm) - 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) + write_pd(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout) def read_from_pd_page( self, @@ -849,254 +276,6 @@ def read_from_pd_page( req_idx: int, pd_page_index: int, state_token_len: int, - layout: StateTransferLayout, + layout, ) -> 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" - self.runtime.clear_history(req_idx) - fields = self._fields( - self.conv_state[:, req_idx, :, : self.config.conv_width], - self.ssm_state[:, req_idx], - ) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - 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]) - - @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")), - } - - -from typing import NamedTuple - - -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode - - -def _find_replay_state(att_state): - manager = getattr(att_state.infer_state, "att_state_manager", None) - if manager is None: - return None - role = att_state.infer_state.att_state_role - matches = [state for state in manager.get_bounded_attention_states(role) if isinstance(state, ReplayKDAState)] - if not matches: - return None - if len(matches) != 1: - raise ValueError(f"expected one ReplayKDAState for {role}, found {len(matches)}") - return matches[0] - - -def prepare_replay_kda_prefill(att_state) -> bool: - """Replay prefill 先折叠 history，并用单槽 SSM 索引。""" - - state = _find_replay_state(att_state) - if state is None: - return False - reqs = att_state.infer_state.b_req_idx - state.merge_accepted_updates(reqs) - att_state.b_conv_buffer_idx = reqs - att_state.b_ssm_buffer_idx = reqs - att_state._replay_state = state - return True - - -def prepare_replay_kda_decode(att_state) -> bool: - """准备 KDA replay decode 的 conv/MTP 索引，并折叠本轮 history。""" - - state = _find_replay_state(att_state) - if state is None: - return False - draft_step = att_state.backend.mtp_step - infer = att_state.infer_state - if draft_step == 0: - att_state.b_conv_buffer_idx = infer.b_req_idx - att_state.b_ssm_buffer_idx = infer.b_req_idx - att_state.b1_mtp_cu_q_seq_len = None - att_state.b_num_accepted_tokens = None - else: - if att_state.backend.uses_dynamic_spec_verify_layout(): - from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, - ) - - ( - att_state.b1_mtp_cu_q_seq_len, - att_state.b_conv_buffer_idx, - att_state.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=state.mtp_state_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - else: - width = draft_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "KDA fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - att_state.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - att_state.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - att_state.b_num_accepted_tokens = state.mtp_state_index[att_state.b_conv_buffer_idx] + 1 - att_state.b_ssm_buffer_idx = att_state.b_conv_buffer_idx - state.prepare_for_decode(att_state.b_conv_buffer_idx, att_state.b1_mtp_cu_q_seq_len) - att_state._replay_state = state - return True + read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index 6a40a05183..f9b40b033a 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -15,42 +15,52 @@ } +def _replay_backend(native_path: str) -> str: + if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": + return native_path + replay_paths = { + "linear.triton:TritonLinearAttBackend": "linear.replay_gdn:ReplayTritonLinearAttBackend", + "linear.flashqla:FlashQlaLinearAttBackend": "linear.replay_gdn:ReplayFlashQlaLinearAttBackend", + } + return replay_paths[native_path] + + def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): args = get_env_start_args() backend_str = _get_backend_str(args.llm_prefill_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_prefill_att_backend_classes[backend_str]) + return load_backend_class(_replay_backend(linear_prefill_att_backend_classes[backend_str])) for backend_name in priority_list: backend_class = linear_prefill_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear prefill attention backend: triton.") - return load_backend_class(backend_class) + return load_backend_class(_replay_backend(backend_class)) if validate(backend_name): logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return load_backend_class(_replay_backend(backend_class)) logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return load_backend_class(_replay_backend("linear.triton:TritonLinearAttBackend")) def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): args = get_env_start_args() backend_str = _get_backend_str(args.llm_decode_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_decode_att_backend_classes[backend_str]) + return load_backend_class(_replay_backend(linear_decode_att_backend_classes[backend_str])) for backend_name in priority_list: backend_class = linear_decode_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear decode attention backend: triton.") - return load_backend_class(backend_class) + return load_backend_class(_replay_backend(backend_class)) if validate(backend_name): logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return load_backend_class(_replay_backend(backend_class)) logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return load_backend_class(_replay_backend("linear.triton:TritonLinearAttBackend")) def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 07d258bd24..986214568f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -114,11 +114,6 @@ class LinearAttPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - self._replay_state = None - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import prepare_replay_gdn_prefill - - if prepare_replay_gdn_prefill(self): - return backend: LinearAttBackend = self.backend mtp_step = backend.mtp_step # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 @@ -149,7 +144,7 @@ def prefill_att( model_role=self.infer_state.att_state_role, state_type_filter="gdn", logical_layer_id=layer_num, - state_class=type(self._replay_state) if self._replay_state is not None else GDNState, + state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 @@ -219,11 +214,6 @@ class LinearAttDecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def init_state(self): - self._replay_state = None - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import prepare_replay_gdn_decode - - if prepare_replay_gdn_decode(self): - return 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() @@ -313,7 +303,7 @@ def decode_att( model_role=self.infer_state.att_state_role, state_type_filter="gdn", logical_layer_id=layer_num, - state_class=type(self._replay_state) if self._replay_state is not None else GDNState, + state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) @@ -327,7 +317,6 @@ def decode_att( b, self.infer_state, layer_weight, - layer_num=layer_num, ) else: # 非 MTP 模式下，使用线性层 decode 状态。 @@ -340,7 +329,6 @@ def decode_att( b, self.infer_state, layer_weight, - layer_num=layer_num, ) return core_attn_out, z @@ -354,7 +342,6 @@ def _gdn_decode_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", - layer_num: int = 0, ): backend: LinearAttBackend = self.backend @@ -377,20 +364,6 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - if self._replay_state is not None: - core_attn_out = self._replay_state.replay_forward( - layer_num, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, - ) - return core_attn_out, z core_attn_out, _ = fused_recurrent_gated_delta_rule( q=query, k=key, @@ -415,7 +388,6 @@ def _gdn_mtp_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", - layer_num: int = 0, ): from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, @@ -437,19 +409,6 @@ def _gdn_mtp_kernel( ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - if self._replay_state is not None: - return self._replay_state.replay_forward( - layer_num, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - cu_seqlens_q, - ) 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 diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index cf443b78bc..cd2f8afb39 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -64,11 +64,6 @@ class KDAPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - self._replay_state = None - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import prepare_replay_kda_prefill - - if prepare_replay_kda_prefill(self): - return 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) @@ -93,7 +88,7 @@ def prefill_att( model_role=self.infer_state.att_state_role, state_type_filter="kda", logical_layer_id=layer_num, - state_class=type(self._replay_state) if self._replay_state is not None else KDAState, + state_class=KDAState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] @@ -143,11 +138,6 @@ class KDADecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def init_state(self): - self._replay_state = None - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import prepare_replay_kda_decode - - if prepare_replay_kda_decode(self): - return mtp_step = self.backend.mtp_step if mtp_step == 0: self._init_normal_decode_state() @@ -236,7 +226,7 @@ def decode_att( model_role=self.infer_state.att_state_role, state_type_filter="kda", logical_layer_id=layer_num, - state_class=type(self._replay_state) if self._replay_state is not None else KDAState, + 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) @@ -256,19 +246,6 @@ 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) - if self._replay_state is not None: - return self._replay_state.replay_forward( - layer_num, - q, - k, - v, - raw_gate.reshape(-1, backend.tp_hidden_size), - raw_beta.reshape(-1, backend.tp_num_heads), - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, - ) output, _ = fused_recurrent_kda( q=q, k=k, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py new file mode 100644 index 0000000000..a5331a6de7 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -0,0 +1,268 @@ +"""GDN replay 的线性 attention backend。 + +Native backend 仍走 fused recurrent。这里单独准备单槽 history，并调用 ReplayGDNState。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( + ReplayGDNState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( + fused_gdn_gating, +) +from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( + conv_pack_gdn_decode_inputs, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def _replay_state(att_state, layer_num: int) -> ReplayGDNState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=ReplayGDNState, + ) + + +def _single_replay_state(att_state) -> ReplayGDNState: + role = att_state.infer_state.att_state_role + matches = [ + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayGDNState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}") + return matches[0] + + +class _ReplayGDNBackend: + """把 prefill/decode 状态换成 replay 实现，维度和 prefill kernel 仍来自具体 backend。""" + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +class ReplayTritonLinearAttBackend(_ReplayGDNBackend, TritonLinearAttBackend): + pass + + +class ReplayFlashQlaLinearAttBackend(_ReplayGDNBackend, FlashQlaLinearAttBackend): + pass + + +@dataclasses.dataclass +class ReplayGDNPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + _single_replay_state(self).merge_accepted_updates(reqs) + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill and att_control.linear_att_prefill_dict is not None + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: LinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = backend._split_qkvzba(params["mixed_qkvzba"]) + return self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight), z + + def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): + backend: LinearAttBackend = self.backend + g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.linear_conv1d.mm_param.weight, + bias=layer_weight.linear_conv1d.bias, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ).transpose(0, 1) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + core_attn_out, last_state = backend.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[self.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[self.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + +@dataclasses.dataclass +class ReplayGDNDecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + state = _single_replay_state(self) + draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) + self._init_indices(draft_step, state.mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + state.prepare_for_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + + def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if draft_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = draft_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "GDN replay decode requires batch_size divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: LinearAttBackend = self.backend + mixed_qkv, z, b, a = backend._split_qkvzba(params["mixed_qkvzba"]) + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) + return core, z + return self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + + def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): + backend: LinearAttBackend = self.backend + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + layer_weight.linear_conv1d.bias, + self.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + output = state.replay_forward( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + return output, z + + def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): + backend: LinearAttBackend = self.backend + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + mtp_step=draft_step, + bias=layer_weight.linear_conv1d.bias, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) + return state.replay_forward( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py new file mode 100644 index 0000000000..94b0b0a4d3 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -0,0 +1,226 @@ +"""KDA replay 的线性 attention backend。 + +Native backend 仍走 fused recurrent。这里单独准备单槽 history，并调用 ReplayKDAState。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( + chunk_kda_with_fused_gate, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def _replay_state(att_state, layer_num: int) -> ReplayKDAState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=ReplayKDAState, + ) + + +def _single_replay_state(att_state) -> ReplayKDAState: + role = att_state.infer_state.att_state_role + matches = [ + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayKDAState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}") + return matches[0] + + +class ReplayKDALinearAttBackend(KDALinearAttBackend): + """维度配置沿用 KDALinearAttBackend，计算状态换成 replay。""" + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayKDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +@dataclasses.dataclass +class ReplayKDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + _single_replay_state(self).merge_accepted_updates(reqs) + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: KDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + conv_states = conv_states[..., : backend.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + params["mixed_qkv"].transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) + k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) + v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].view(1, -1, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(1, -1, backend.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states[self.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +@dataclasses.dataclass +class ReplayKDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + state = _single_replay_state(self) + self._init_indices(self.backend.mtp_step, state.mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + state.prepare_for_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + + def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if mtp_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = mtp_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "KDA replay decode requires batch_size divisible by mtp_step + 1, " + f"got batch_size={batch_size}, mtp_step={mtp_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: KDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + params["mixed_qkv"], + conv_states, + layer_weight.get_merged_kda_conv_weight(), + **conv_kwargs, + ) + q, k, v = backend.split_qkv(mixed_qkv) + shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) + q = q.view(*shape, backend.tp_num_heads, backend.head_dim) + k = k.view(*shape, backend.tp_num_heads, backend.head_dim) + v = v.view(*shape, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].view(*shape, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) + return state.replay_forward( + layer_num, + q, + k, + v, + raw_gate.reshape(-1, backend.tp_hidden_size), + raw_beta.reshape(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index a2a7d10ef7..f1d0486966 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -158,7 +158,10 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + from lightllm.common.attention_state.full_att import ( + PackedMLALogicalConfig, + create_packed_mla_state, + ) from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -192,7 +195,9 @@ def prepare_attention_states(self) -> None: Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) def _init_att_backend(self): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + ) backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend @@ -201,7 +206,15 @@ def _init_att_backend(self): self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): - self.prefill_att_backend1 = KDALinearAttBackend(model=self) + if getattr(self.args, "ssm_state_mode", "native") == "replay": + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDALinearAttBackend, + ) + + backend_cls = ReplayKDALinearAttBackend + else: + backend_cls = KDALinearAttBackend + self.prefill_att_backend1 = backend_cls(model=self) self.decode_att_backend1 = self.prefill_att_backend1 def _init_custom(self): diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index b531a025b7..d6ee8cd54e 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -20,10 +20,10 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime return ReplayGDNRuntime(*args, **kwargs) 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 index bc51f0fc4a..0c0bfac0e6 100644 --- 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 @@ -8,17 +8,17 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime return ReplayGDNRuntime(*args, **kwargs) from types import SimpleNamespace from lightllm.common.attention_state.commit import StateAcceptance -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix 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 index 236d9bff8a..49da74458c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -5,10 +5,10 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime return ReplayGDNRuntime(*args, **kwargs) 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 index 3f6795b0a6..fc74c1224f 100644 --- 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 @@ -6,10 +6,10 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime return ReplayGDNRuntime(*args, **kwargs) 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 index 4162424509..4dd9d61d5f 100644 --- 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 @@ -8,10 +8,10 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime return ReplayGDNRuntime(*args, **kwargs) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index f7223c957e..5dce20a614 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -8,7 +8,8 @@ from lightllm.common.attention_state.bounded_att import KDAState from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNRuntime, ReplayGDNState +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index f1a7a151f9..24ccc83819 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -5,7 +5,9 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( + fused_recurrent_kda, +) from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -23,7 +25,11 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): - cu = torch.tensor([0, lengths[0], sum(lengths), sum(lengths)], device="cuda", dtype=torch.int32) + cu = torch.tensor( + [0, lengths[0], sum(lengths), sum(lengths)], + device="cuda", + dtype=torch.int32, + ) tokens = sum(lengths) q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -100,7 +106,12 @@ def prefill(values, destination, tail): prefill(raw[:prefix], packed, ring) # Alternate expanded verification, rewound single-token drafting, and re-verification. - for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + for start, count in [ + (prefix, width), + (prefix + 1, 1), + (prefix + 2, 1), + (prefix + 1, width), + ]: raw[start : start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) compress_pools( @@ -117,7 +128,11 @@ def prefill(values, destination, tail): b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) - prefill(raw[: start + count], expected, torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16)) + prefill( + raw[: start + count], + expected, + torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16), + ) closing = torch.arange(start, start + count, device="cuda") closing = closing[(closing + 1) % 4 == 0] assert torch.equal(packed[closing], expected[closing]) @@ -131,8 +146,15 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDADecodeAttState, + ) + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, + ) + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_update, + ) torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 @@ -234,7 +256,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState + state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() actual = state.decode_att( None, @@ -243,15 +266,27 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod AttControl( linear_att_decode=True, linear_att_decode_dict=dict( - layer_weight=weight, layer_num=0, mixed_qkv=mixed, raw_gate=gate, raw_beta=beta + layer_weight=weight, + layer_num=0, + mixed_qkv=mixed, + raw_gate=gate, + raw_beta=beta, ), ), ) - torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) + torch.testing.assert_close( + actual.reshape(tokens, heads, dim)[:row], + torch.cat(expected), + atol=2e-3, + rtol=1e-2, + ) if mode == "native": checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] torch.testing.assert_close( - selected.ssm_state[0, :checked_slots], expected_ssm[:checked_slots], atol=2e-3, rtol=1e-2 + selected.ssm_state[0, :checked_slots], + expected_ssm[:checked_slots], + atol=2e-3, + rtol=1e-2, ) else: ids = torch.tensor(requests, dtype=torch.int32, device="cuda") From 794882421555df5feea35a82055bc6dc718cb5cb Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 02:08:01 +0000 Subject: [PATCH 08/54] refactor: make replay KDA and GDN backends standalone Replay inference no longer subclasses the native linear backends. Each algorithm owns its entry point and configures its own prefill kernel. Co-authored-by: Cursor --- .../attention/create_linear_utils.py | 30 ++-- .../basemodel/attention/linear/replay_gdn.py | 132 +++++++++++++++--- .../basemodel/attention/linear/replay_kda.py | 29 +++- .../basemodel/attention/linear/test_gdn.py | 28 ++++ 4 files changed, 177 insertions(+), 42 deletions(-) diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index f9b40b033a..ffd5e7eeb9 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -15,52 +15,52 @@ } -def _replay_backend(native_path: str) -> str: +def _linear_backend_class(native_path: str, backend_name: str): + """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" + if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": - return native_path - replay_paths = { - "linear.triton:TritonLinearAttBackend": "linear.replay_gdn:ReplayTritonLinearAttBackend", - "linear.flashqla:FlashQlaLinearAttBackend": "linear.replay_gdn:ReplayFlashQlaLinearAttBackend", - } - return replay_paths[native_path] + return load_backend_class(native_path) + from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class + + return replay_gdn_backend_class(backend_name) def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): args = get_env_start_args() backend_str = _get_backend_str(args.llm_prefill_att_backend, index) if backend_str != "auto": - return load_backend_class(_replay_backend(linear_prefill_att_backend_classes[backend_str])) + return _linear_backend_class(linear_prefill_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_prefill_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear prefill attention backend: triton.") - return load_backend_class(_replay_backend(backend_class)) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return load_backend_class(_replay_backend(backend_class)) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return load_backend_class(_replay_backend("linear.triton:TritonLinearAttBackend")) + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): args = get_env_start_args() backend_str = _get_backend_str(args.llm_decode_att_backend, index) if backend_str != "auto": - return load_backend_class(_replay_backend(linear_decode_att_backend_classes[backend_str])) + return _linear_backend_class(linear_decode_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_decode_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear decode attention backend: triton.") - return load_backend_class(_replay_backend(backend_class)) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return load_backend_class(_replay_backend(backend_class)) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return load_backend_class(_replay_backend("linear.triton:TritonLinearAttBackend")) + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index a5331a6de7..eac3b0893b 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -1,6 +1,7 @@ -"""GDN replay 的线性 attention backend。 +"""GDN replay 的推理入口。 -Native backend 仍走 fused recurrent。这里单独准备单槽 history，并调用 ReplayGDNState。 +Triton 和 FlashQla 各有一个 backend。它们只差 prefill chunk kernel， +decode 都走 ReplayGDNState。两个类都直接实现 backend 接口，不继承 native GDN。 """ from __future__ import annotations @@ -16,12 +17,10 @@ ) from lightllm.common.basemodel.attention.base_att import ( AttControl, + BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( causal_conv1d_fn, ) @@ -37,11 +36,94 @@ from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) +from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo +def replay_gdn_backend_class(name: str): + """按已选中的 linear kernel 名字返回对应的 replay backend。""" + + classes = { + "triton": ReplayGDNTritonAttBackend, + "flashqla": ReplayGDNFlashQlaAttBackend, + } + try: + return classes[name] + except KeyError: + raise ValueError(f"replay GDN has no backend named {name!r}") from None + + +def _triton_prefill_kernel(): + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + chunk_gated_delta_rule, + ) + + return chunk_gated_delta_rule + + +def _flashqla_prefill_kernel(): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +def _configure_gdn(backend, model, prefill_kernel) -> None: + """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" + + config = model.config + tp_world_size = model.tp_world_size_ + backend.mtp_step = get_env_start_args().mtp_step + backend.num_v_heads = config["linear_num_value_heads"] + backend.num_k_heads = config["linear_num_key_heads"] + backend.head_k_dim = config["linear_key_head_dim"] + backend.head_v_dim = config["linear_value_head_dim"] + backend.key_dim = backend.head_k_dim * backend.num_k_heads + backend.value_dim = backend.head_v_dim * backend.num_v_heads + backend.conv_kernel_dim = config["linear_conv_kernel_dim"] + backend.activation = config["hidden_act"] + backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size + backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size + backend.tp_num_k_heads = backend.num_k_heads // tp_world_size + backend.tp_num_v_heads = backend.num_v_heads // tp_world_size + backend.tp_key_dim = backend.key_dim // tp_world_size + backend.tp_value_dim = backend.value_dim // tp_world_size + assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" + backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads + ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + backend.prefill_kernel = prefill_kernel + + +def _split_qkvzba(backend, mixed_qkvzba): + qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim + z_end = qkv_dim + backend.tp_value_dim + b_end = z_end + backend.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return mixed_qkv, z, b, a + + +def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + if decode: + batch_size = mixed_qkv.shape[0] + leading = (batch_size, 1) + else: + leading = (1, query.shape[0]) + query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) + return query, key, value + + def _replay_state(att_state, layer_num: int) -> ReplayGDNState: return att_state.infer_state.att_state_manager.get_state( model_role=att_state.infer_state.att_state_role, @@ -63,8 +145,12 @@ def _single_replay_state(att_state) -> ReplayGDNState: return matches[0] -class _ReplayGDNBackend: - """把 prefill/decode 状态换成 replay 实现，维度和 prefill kernel 仍来自具体 backend。""" +class ReplayGDNTritonAttBackend(BaseAttBackend): + """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + _configure_gdn(self, model, _triton_prefill_kernel()) def create_att_prefill_state(self, infer_state: "InferStateInfo"): return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -73,12 +159,18 @@ def create_att_decode_state(self, infer_state: "InferStateInfo"): return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) -class ReplayTritonLinearAttBackend(_ReplayGDNBackend, TritonLinearAttBackend): - pass +class ReplayGDNFlashQlaAttBackend(BaseAttBackend): + """GDN replay 的 FlashQla 入口。Prefill 使用 flash_qla chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + _configure_gdn(self, model, _flashqla_prefill_kernel()) + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) -class ReplayFlashQlaLinearAttBackend(_ReplayGDNBackend, FlashQlaLinearAttBackend): - pass + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @dataclasses.dataclass @@ -104,16 +196,16 @@ def prefill_att( params = att_control.linear_att_prefill_dict layer_num = params["layer_num"] layer_weight = params["layer_weight"] - backend: LinearAttBackend = self.backend + backend = self.backend state = _replay_state(self, layer_num) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) if backend.mtp_step > 0: conv_states = conv_states[:, :, : -backend.mtp_step] - mixed_qkv, z, b, a = backend._split_qkvzba(params["mixed_qkvzba"]) + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) return self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight), z def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): - backend: LinearAttBackend = self.backend + backend = self.backend g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -125,7 +217,7 @@ def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): conv_states=conv_states, activation=backend.activation, ).transpose(0, 1) - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv) core_attn_out, last_state = backend.prefill_kernel( q=query, k=key, @@ -198,8 +290,8 @@ def decode_att( params = att_control.linear_att_decode_dict layer_num = params["layer_num"] layer_weight = params["layer_weight"] - backend: LinearAttBackend = self.backend - mixed_qkv, z, b, a = backend._split_qkvzba(params["mixed_qkvzba"]) + backend = self.backend + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) state = _replay_state(self, layer_num) conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) @@ -209,7 +301,7 @@ def decode_att( return self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): - backend: LinearAttBackend = self.backend + backend = self.backend query, key, value, z, a, b = conv_pack_gdn_decode_inputs( mixed_qkv, z, @@ -241,7 +333,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat return output, z def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): - backend: LinearAttBackend = self.backend + backend = self.backend mixed_qkv = causal_conv1d_update_mtp( mixed_qkv, conv_states, @@ -253,7 +345,7 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step num_accepted_tokens=self.b_num_accepted_tokens, query_start_loc=self.b1_mtp_cu_q_seq_len, ) - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) + query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) return state.replay_forward( layer_num, query, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 94b0b0a4d3..484b9629d0 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -1,6 +1,6 @@ -"""KDA replay 的线性 attention backend。 +"""KDA replay 的推理入口。 -Native backend 仍走 fused recurrent。这里单独准备单槽 history，并调用 ReplayKDAState。 +这个 backend 只服务 ReplayKDAState。维度在这里自己读取，不继承 native KDA backend。 """ from __future__ import annotations @@ -16,10 +16,10 @@ ) from lightllm.common.basemodel.attention.base_att import ( AttControl, + BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, @@ -59,8 +59,23 @@ def _single_replay_state(att_state) -> ReplayKDAState: return matches[0] -class ReplayKDALinearAttBackend(KDALinearAttBackend): - """维度配置沿用 KDALinearAttBackend，计算状态换成 replay。""" +class ReplayKDALinearAttBackend(BaseAttBackend): + """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" + + def __init__(self, model): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) def create_att_prefill_state(self, infer_state: "InferStateInfo"): return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -92,7 +107,7 @@ def prefill_att( params = att_control.linear_att_prefill_dict layer_num = params["layer_num"] layer_weight = params["layer_weight"] - backend: KDALinearAttBackend = self.backend + backend: ReplayKDALinearAttBackend = self.backend state = _replay_state(self, layer_num) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] @@ -187,7 +202,7 @@ def decode_att( params = att_control.linear_att_decode_dict layer_num = params["layer_num"] layer_weight = params["layer_weight"] - backend: KDALinearAttBackend = self.backend + backend: ReplayKDALinearAttBackend = self.backend state = _replay_state(self, layer_num) conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..b61b9c02b3 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -177,6 +177,34 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear assert validate_calls == ["flashqla"] +def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace( + llm_prefill_att_backend=["fa3", "flashqla"], + llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="replay", + ), + ) + from lightllm.common.basemodel.attention.base_att import BaseAttBackend + from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_gdn import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, + ) + + prefill_backend = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + + assert prefill_backend is ReplayGDNFlashQlaAttBackend + assert decode_backend is ReplayGDNTritonAttBackend + assert prefill_backend.__bases__ == (BaseAttBackend,) + assert decode_backend.__bases__ == (BaseAttBackend,) + assert not issubclass(prefill_backend, LinearAttBackend) + assert not issubclass(decode_backend, LinearAttBackend) + + def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, From cb6f6f8b9ec27a430323d2bf3f44f2e8430de127 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 02:26:47 +0000 Subject: [PATCH 09/54] refactor: move replay decode onto GDN and KDA step modules History runtimes keep request storage, materialization, and copies. The inference step owns kernel selection and the per-layer replay. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_runtime.py | 99 +--------------- .../bounded_att/gdn/replay_gdn_state.py | 42 ++----- .../bounded_att/kda/replay_kda_runtime.py | 99 +--------------- .../bounded_att/kda/replay_kda_state.py | 42 ++----- .../basemodel/attention/linear/replay_gdn.py | 15 ++- .../attention/linear/replay_gdn_step.py | 110 ++++++++++++++++++ .../basemodel/attention/linear/replay_kda.py | 10 +- .../attention/linear/replay_kda_step.py | 110 ++++++++++++++++++ .../triton_kernel/linear_att/ssm_autotune.py | 12 +- test/benchmark/kernels/benchmark_kda_ssm.py | 4 +- .../linear_att/test_acceptance_fusion.py | 3 +- .../linear_att/test_replayssm.py | 15 +-- .../linear_att/test_replayssm_execution.py | 13 ++- .../linear_att/test_ssm_autotune_memory.py | 2 +- .../common/test_linear_state_contract.py | 3 +- 15 files changed, 291 insertions(+), 288 deletions(-) create mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn_step.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_kda_step.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py index 828d535c14..f5e3c8d934 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py @@ -1,4 +1,4 @@ -"""GDN replay history。只负责 ring buffer、重放 kernel 和 accepted SSM 物化。""" +"""GDN replay history。只保存 ring buffer、cursor，以及 accepted SSM 的物化和拷贝。""" from __future__ import annotations @@ -13,10 +13,6 @@ _prepare_fold, _fold_history, _materialize, - _replay, -) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import ( - resolve_run_config, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( commit_ssm_prefix, @@ -30,7 +26,7 @@ def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Te class ReplayGDNRuntime: - """以短 history 重放 decode 的 SSM 执行器。 + """请求槽上的 GDN replay history。 ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode @@ -42,7 +38,7 @@ class ReplayGDNRuntime: @property def run_config(self) -> SSMRunConfig | None: - """返回最近一次 forward 使用的 kernel 配置。""" + """返回物化和 autotune 使用的 kernel 配置。""" return self._run_config @@ -286,95 +282,6 @@ def export_accepted_state( output_strides=strides, ) - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> SSMRunConfig: - """按当前 kernel 形状选择或返回固定的 GDN replay 调优配置。""" - - return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig]: - """执行一层 GDN replay decode，并返回输出与本轮 run config。""" - - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - 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 == "precompute" - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor - # 运算，无法定义静态 shape。 - 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, - self.kda, - self.gate_lower_bound, - 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: torch.Tensor) -> None: """非 MTP decode 直接把单 token 写入当前 history。""" self.accept_updates(reqs) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 4c1c9ac2e7..0faf04ec49 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -2,6 +2,7 @@ 直接实现状态接口，不继承 native GDNState。history 在 replay_gdn_runtime， 页布局在 replay_gdn_checkpoint，已接受快照在 replay_gdn_snapshot。 +一层 decode 重放在 replay_gdn_step。 """ from __future__ import annotations @@ -137,8 +138,13 @@ def initialize_page_pools( def contains_logical_layer(self, logical_layer_id: int) -> bool: return logical_layer_id in self._logical_to_storage_layer + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int): - index = self._logical_to_storage_layer[logical_layer_id] + index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] @deepcopy_result @@ -214,40 +220,6 @@ def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """Prefill 前把已接受 history 折回 canonical state。""" self.runtime.merge_accepted_updates(reqs) - def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" - self.runtime.prepare_decode(reqs, cu_seqlens) - if cu_seqlens is None: - self.runtime.accept_single_token(reqs) - - def replay_forward( - self, - logical_layer: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> torch.Tensor: - """执行一层 replay decode。""" - output, _ = self.runtime.forward( - self._logical_to_storage_layer[logical_layer], - q, - k, - v, - a, - b, - a_log, - bias, - reqs, - cu_seqlens, - ) - return output - def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py index 53c2f75447..ae5a734dba 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py @@ -1,4 +1,4 @@ -"""KDA replay history。只负责 ring buffer、重放 kernel 和 accepted SSM 物化。""" +"""KDA replay history。只保存 ring buffer、cursor，以及 accepted SSM 的物化和拷贝。""" from __future__ import annotations @@ -13,10 +13,6 @@ _prepare_fold, _fold_history, _materialize, - _replay, -) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import ( - resolve_run_config, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( commit_ssm_prefix, @@ -30,7 +26,7 @@ def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Te class ReplayKDARuntime: - """以短 history 重放 decode 的 SSM 执行器。 + """请求槽上的 KDA replay history。 ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode @@ -42,7 +38,7 @@ class ReplayKDARuntime: @property def run_config(self) -> SSMRunConfig | None: - """返回最近一次 forward 使用的 kernel 配置。""" + """返回物化和 autotune 使用的 kernel 配置。""" return self._run_config @@ -286,95 +282,6 @@ def export_accepted_state( output_strides=strides, ) - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> SSMRunConfig: - """按当前 kernel 形状选择或返回固定的 KDA replay 调优配置。""" - - return resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, - *, - state_indices: torch.Tensor | None = None, - accepted_token_counts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, SSMRunConfig]: - """执行一层 KDA replay decode，并返回输出与本轮 run config。""" - - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - 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 == "precompute" - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor - # 运算，无法定义静态 shape。 - 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, - self.kda, - self.gate_lower_bound, - 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: torch.Tensor) -> None: """非 MTP decode 直接把单 token 写入当前 history。""" self.accept_updates(reqs) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 0e877746a8..5b41edb4fc 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -2,6 +2,7 @@ 直接实现状态接口，不继承 native KDAState。history 在 replay_kda_runtime， 页布局在 replay_kda_checkpoint，已接受快照在 replay_kda_snapshot。 +一层 decode 重放在 replay_kda_step。 """ from __future__ import annotations @@ -137,8 +138,13 @@ def initialize_page_pools( def contains_logical_layer(self, logical_layer_id: int) -> bool: return logical_layer_id in self._logical_to_storage_layer + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int): - index = self._logical_to_storage_layer[logical_layer_id] + index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] @deepcopy_result @@ -214,40 +220,6 @@ def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """Prefill 前把已接受 history 折回 canonical state。""" self.runtime.merge_accepted_updates(reqs) - def prepare_for_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" - self.runtime.prepare_decode(reqs, cu_seqlens) - if cu_seqlens is None: - self.runtime.accept_single_token(reqs) - - def replay_forward( - self, - logical_layer: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - ) -> torch.Tensor: - """执行一层 replay decode。""" - output, _ = self.runtime.forward( - self._logical_to_storage_layer[logical_layer], - q, - k, - v, - a, - b, - a_log, - bias, - reqs, - cu_seqlens, - ) - return output - def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index eac3b0893b..e5ebe6d83a 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -38,6 +38,8 @@ ) from lightllm.utils.envs_utils import get_env_start_args +from .replay_gdn_step import begin_decode, replay + if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -245,7 +247,7 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) self._init_indices(draft_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - state.prepare_for_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + begin_decode(state.runtime, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state @@ -318,8 +320,9 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat backend.tp_num_v_heads, backend.head_v_dim, ) - output = state.replay_forward( - layer_num, + output, _ = replay( + state.runtime, + state.storage_layer(layer_num), query, key, value, @@ -346,8 +349,9 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step query_start_loc=self.b1_mtp_cu_q_seq_len, ) query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) - return state.replay_forward( - layer_num, + output, _ = replay( + state.runtime, + state.storage_layer(layer_num), query, key, value, @@ -358,3 +362,4 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) + return output diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py new file mode 100644 index 0000000000..db887cb923 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py @@ -0,0 +1,110 @@ +"""GDN replay 的一层 decode。 + +History 张量和 cursor 留在 ReplayGDNRuntime。这里只选择 kernel 配置并发起重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from lightllm.common.ssm_types import SSMRunConfig + + +def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: + """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" + + history.prepare_decode(reqs, cu_seqlens) + if cu_seqlens is None: + history.accept_single_token(reqs) + + +def run_config( + history, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, +) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 GDN replay 调优配置。""" + + return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, +) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 GDN replay decode，并返回输出与本轮 run config。""" + + if history.kda: + raise ValueError("GDN replay step does not use the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.state.shape[-3:] + assert history.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 = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.forward_cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + False, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 484b9629d0..d8325fcf20 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -34,6 +34,8 @@ build_dynamic_mtp_linear_att_state_params, ) +from .replay_kda_step import begin_decode, replay + if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -157,7 +159,7 @@ def init_state(self): state = _single_replay_state(self) self._init_indices(self.backend.mtp_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - state.prepare_for_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + begin_decode(state.runtime, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state @@ -227,8 +229,9 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = params["raw_gate"].view(*shape, backend.tp_hidden_size) raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) - return state.replay_forward( - layer_num, + output, _ = replay( + state.runtime, + state.storage_layer(layer_num), q, k, v, @@ -239,3 +242,4 @@ def decode_att( self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) + return output diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_step.py b/lightllm/common/basemodel/attention/linear/replay_kda_step.py new file mode 100644 index 0000000000..5fcfa29d45 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda_step.py @@ -0,0 +1,110 @@ +"""KDA replay 的一层 decode。 + +History 张量和 cursor 留在 ReplayKDARuntime。这里只选择 kernel 配置并发起重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from lightllm.common.ssm_types import SSMRunConfig + + +def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: + """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" + + history.prepare_decode(reqs, cu_seqlens) + if cu_seqlens is None: + history.accept_single_token(reqs) + + +def run_config( + history, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, +) -> SSMRunConfig: + """按当前 kernel 形状选择或返回固定的 KDA replay 调优配置。""" + + return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: SSMRunConfig | None = None, +) -> tuple[torch.Tensor, SSMRunConfig]: + """执行一层 KDA replay decode，并返回输出与本轮 run config。""" + + if not history.kda: + raise ValueError("KDA replay step requires the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.state.shape[-3:] + assert history.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 = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.forward_cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + True, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index c4e78a2b65..2556e1681b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -129,6 +129,16 @@ def rebuild_inputs( ) +def replay_layer(history, *args, **kwargs): + """按 history 所属算法调用对应的 replay decode。""" + + if history.kda: + from lightllm.common.basemodel.attention.linear.replay_kda_step import replay + else: + from lightllm.common.basemodel.attention.linear.replay_gdn_step import replay + return replay(history, *args, **kwargs) + + @autotune( kernel_name="ssm_update_cycle:v1", configs_gen_func=get_configs, @@ -153,7 +163,7 @@ def select_config( rounds = ssm_backend.capacity // min(ssm_backend.verify_width, tokens) + 1 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) + replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) ssm_backend.accept_updates(reqs, accepted) return run_config diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index d6ee8cd54e..14ea97b744 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -28,7 +28,7 @@ def ReplaySSMExecutor(*args, **kwargs): return ReplayGDNRuntime(*args, **kwargs) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer def workload(args, batch, width, dtype, mode, acceptance): @@ -84,7 +84,7 @@ def cycle(): 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) + replay_layer(ssm_backend, 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, ReplaySSMExecutor): 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 index 0c0bfac0e6..f71243a216 100644 --- 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 @@ -20,6 +20,7 @@ def ReplaySSMExecutor(*args, **kwargs): from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer def make(n, width): @@ -181,7 +182,7 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput 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)] + outputs = [replay_layer(ssm_backend, 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 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 index 49da74458c..dccd92afba 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -13,7 +13,7 @@ def ReplaySSMExecutor(*args, **kwargs): return ReplayGDNRuntime(*args, **kwargs) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -73,7 +73,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state # an overflowed history 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] + actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] expected = [] for token in range(2): output, expected_state = reference( @@ -178,7 +178,7 @@ def test_replay_acceptance_flush_and_graph( 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)] + out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] ssm_backend.accept_updates(reqs, accepted) return out @@ -336,7 +336,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d 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]) + outputs.append(replay_layer(ssm_backend, 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) @@ -374,7 +374,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): 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]) + outputs.append(replay_layer(ssm_backend, 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) @@ -401,7 +401,7 @@ def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, 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)] + return [replay_layer(ssm_backend, 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) @@ -636,7 +636,8 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa inline = backend.forward_cursors[:live] % 16 + query_width > 8 assert inline[:folding].eq(not separate).all() assert not inline[folding:].any() - actual = backend.forward( + actual = replay_layer( + backend, 0, q.reshape(1, -1, 1, 32), k.reshape(1, -1, 1, 32), 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 index fc74c1224f..77436ca1ff 100644 --- 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 @@ -14,6 +14,7 @@ def ReplaySSMExecutor(*args, **kwargs): return ReplayGDNRuntime(*args, **kwargs) +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) @@ -57,7 +58,7 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() backend.prepare_decode(ids) - out = backend.forward(0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] + out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] backend.accept_single_token(ids) return out @@ -157,7 +158,7 @@ def metadata(ssm_backend): 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)] + out = [replay_layer(ssm_backend, 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: @@ -250,7 +251,7 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, 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)] + return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] def warm(ssm_backend): ssm_backend.prepare_decode(reqs) @@ -336,7 +337,7 @@ def test_cross_layer_fold_with_state_offset_above_int32(): 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)] + out = [replay_layer(ssm_backend, 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 @@ -385,7 +386,9 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): 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)] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2) + ] ssm_backend.accept_updates(ids, accepted) return out 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 index 4dd9d61d5f..2ee7477f54 100644 --- 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 @@ -76,5 +76,5 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm 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) + tuning.replay_layer(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) torch.cuda.synchronize() diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 5dce20a614..7e0e6fb870 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -9,6 +9,7 @@ from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime +from lightllm.common.basemodel.attention.linear.replay_gdn_step import run_config as gdn_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -97,7 +98,7 @@ def test_replay_backend_publishes_detached_run_config(): ssm_backend.run_config = tuning ssm_backend._config_is_fixed = True ssm_backend.compute_mode = "inline" - config = ssm_backend.get_run_config(*([None] * 7)) + config = gdn_run_config(ssm_backend, *([None] * 7)) assert isinstance(config, SSMRunConfig) tuning["BV"] = 128 assert config.BV == 32 From 98957a437848debe94eb87c5fac2a9372263e477 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 02:57:07 +0000 Subject: [PATCH 10/54] refactor: keep GDN and KDA replay history on the state The separate runtime objects only duplicated the canonical SSM. History buffers, cursors, and accepted-state materialization now live on the state. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_runtime.py | 332 ---------------- .../bounded_att/gdn/replay_gdn_snapshot.py | 12 +- .../bounded_att/gdn/replay_gdn_state.py | 360 +++++++++++++++++- .../bounded_att/kda/replay_kda_runtime.py | 332 ---------------- .../bounded_att/kda/replay_kda_snapshot.py | 12 +- .../bounded_att/kda/replay_kda_state.py | 360 +++++++++++++++++- .../basemodel/attention/linear/replay_gdn.py | 6 +- .../attention/linear/replay_gdn_step.py | 6 +- .../basemodel/attention/linear/replay_kda.py | 4 +- .../attention/linear/replay_kda_step.py | 6 +- .../triton_kernel/linear_att/ssm_autotune.py | 10 +- test/benchmark/kernels/benchmark_kda_ssm.py | 8 +- .../test_hybrid_layer_mapping.py | 4 +- .../linear_att/test_acceptance_fusion.py | 18 +- .../linear_att/test_replayssm.py | 20 +- .../linear_att/test_replayssm_execution.py | 8 +- .../linear_att/test_ssm_autotune_memory.py | 8 +- .../test_hybrid_req_manager_pd_sharing.py | 9 +- unit_tests/common/test_linear_state.py | 6 +- .../common/test_linear_state_contract.py | 22 +- unit_tests/models/glm5_next/test_cache.py | 2 +- unit_tests/models/glm5_next/test_mtp.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 2 +- 23 files changed, 772 insertions(+), 777 deletions(-) delete mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py delete mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py deleted file mode 100644 index f5e3c8d934..0000000000 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_runtime.py +++ /dev/null @@ -1,332 +0,0 @@ -"""GDN replay history。只保存 ring buffer、cursor，以及 accepted SSM 的物化和拷贝。""" - -from __future__ import annotations - -from collections.abc import Sequence - -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, -) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( - commit_ssm_prefix, -) - - -def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -class ReplayGDNRuntime: - """请求槽上的 GDN replay history。 - - ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 - cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode - 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 - checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 - """ - - _run_config: SSMRunConfig | None = None - - @property - def run_config(self) -> SSMRunConfig | None: - """返回物化和 autotune 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" - - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) - - def __init__( - self, - state: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - *, - num_key_heads: int | None = None, - kda: bool = False, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, - ) -> None: - """创建 GDN replay 的 canonical state、双半区 history 和 cursor。 - - ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 - """ - - assert state.dtype in ( - torch.float32, - torch.bfloat16, - ), "GDN replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - if kda: - raise ValueError("GDN replay runtime does not use the KDA gate") - self.state = state - self.kda = False - 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 - # 保留冗余布局，兼容已有调用方和 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) - # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 - # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 - # 下可以使用更窄的 history buffer；其他形状使用 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), - 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: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - 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, - 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, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor 以保留每个请求本轮已接受的前缀。 - - MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft - 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 - """ - _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: torch.Tensor) -> None: - """把接受的 history 物化回 canonical state，并清空请求 cursor。""" - 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: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - - layers, slots, hv, k, v = self.state.shape - config = self.run_config or SSMRunConfig(32, 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, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回不修改 active state/history 的单请求 canonical 快照。""" - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" - - if not req_indices: - return - shape = (self.state.shape[0], *self.state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) - self._materialize_accepted_state( - upload_indices(req_indices, self.state.device, torch.int32), - upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - - def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 直接把单 token 写入当前 history。""" - self.accept_updates(reqs) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) - - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" - - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - - def export_ssm( - self, - req_indices: list[int], - targets: list[torch.Tensor], - *, - flush: bool = False, - ) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.state.shape[0], *self.state.shape[2:]), - dtype=self.state.dtype, - device=self.state.device, - ) - self._materialize_accepted_state( - torch.tensor([req], dtype=torch.int32, device=self.state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.state[:, req].copy_(snap) - self.clear_history(req) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py index 9eaecd0d2e..f6492f34b9 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py @@ -33,7 +33,7 @@ def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) for req_idx, slot_id in zip(req_indices, page_slot_ids) ] ) - state.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + state.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: @@ -44,7 +44,7 @@ def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_i width = state.config.conv_width conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - state.runtime.clear_history(req_idx) + state.clear_history(req_idx) conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) @@ -64,7 +64,7 @@ def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_boun width = state.config.conv_width conv_copies = [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - state.runtime.clear_history(dst_req) + state.clear_history(dst_req) conv_copies.append( ( accepted_conv(remote_dp_bounded_state, src_req), @@ -72,9 +72,7 @@ def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_boun ) ) copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.runtime.export_ssm( - src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices] - ) + remote_dp_bounded_state.export_ssm(src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices]) if state.mtp_state_index is not None: for dst_req in dst_req_indices: state.mtp_state_index[dst_req].fill_(0) @@ -100,7 +98,7 @@ def read_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, lay """读入基准 Conv/SSM，并清空该请求的 history。""" assert layout.state_name == state.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - state.runtime.clear_history(req_idx) + state.clear_history(req_idx) fields = tensor_fields( state.conv_state[:, req_idx, :, : state.config.conv_width], state.ssm_state[:, req_idx], diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 0faf04ec49..7abfbada28 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -1,16 +1,26 @@ """GDN replay 的 bounded attention 状态。 -直接实现状态接口，不继承 native GDNState。history 在 replay_gdn_runtime， -页布局在 replay_gdn_checkpoint，已接受快照在 replay_gdn_snapshot。 -一层 decode 重放在 replay_gdn_step。 +Conv、canonical SSM 和 replay history 都定义在本类。页布局在 replay_gdn_checkpoint， +已接受快照在 replay_gdn_snapshot。一层 decode 重放在 replay_gdn_step。 """ from __future__ import annotations +from collections.abc import Sequence from functools import lru_cache from typing import NamedTuple import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _accept_updates, + _fold_history, + _materialize, + _prepare_fold, +) +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.ssm_types import SSMRunConfig from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState @@ -22,7 +32,6 @@ read_pages_from_cpu, write_pages_to_cpu, ) -from .replay_gdn_runtime import ReplayGDNRuntime from .replay_gdn_snapshot import ( accepted_conv, copy_runtime_from_dp, @@ -45,8 +54,84 @@ def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: return capacity, compute_mode +def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: SSMRunConfig | dict[str, int] | None, +) -> None: + """在 owner 上创建 GDN replay 的双半区 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = False + owner.gate_lower_bound = gate_lower_bound + owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.forward_cursors = torch.zeros_like(owner.cursors) + # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + + class ReplayGDNState(BoundedAttentionState): - """GDN replay 运行态。槽位布局和 history 都在本类，不沿用 native 的 MTP 槽。""" + """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + _run_config: SSMRunConfig | None = None + + @property + def state(self) -> torch.Tensor: + """Replay kernel 读取的 canonical SSM，与 ssm_state 是同一张量。""" + + return self.ssm_state + + @property + def run_config(self) -> SSMRunConfig | None: + """返回物化和 autotune 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) def __init__( self, @@ -96,17 +181,60 @@ def __init__( self.big_pages_pool = None self.small_pages_pool = None capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) - self.runtime = ReplayGDNRuntime( + _bind_history( + self, self.ssm_state, capacity, mtp_step + 1, config.conv_dtype, - num_key_heads=config.num_k_heads, - kda=False, - gate_lower_bound=gate_lower_bound, - compute_mode=compute_mode, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, ) + @classmethod + def from_ssm( + cls, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = False, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> "ReplayGDNState": + """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" + + if kda: + raise ValueError("GDN replay does not use the KDA gate") + owner = cls.__new__(cls) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + @property def ssm_slots_per_req(self) -> int: """每个请求只保留一份 canonical SSM。""" @@ -115,9 +243,9 @@ def ssm_slots_per_req(self) -> int: @property def _ssm_executor(self): - """测试沿用的 history 入口。""" + """History 张量在本对象上。测试沿用这个入口。""" - return self.runtime + return self def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" @@ -188,7 +316,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: for req_idx in req_indices: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() - self.runtime.clear_history(req_idx) + self.clear_history(req_idx) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) @@ -214,15 +342,11 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to """提交已接受前缀，并推进 replay cursor。""" if self.mtp_state_index is None: return - self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """Prefill 前把已接受 history 折回 canonical state。""" - self.runtime.merge_accepted_updates(reqs) + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) + return _AcceptedTensors(accepted_conv(self, req_idx), self.accepted_ssm(req_idx)) def copy_runtime_state_from_dp( self, @@ -251,3 +375,201 @@ def read_from_pd_page( layout, ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.ssm_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, + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.ssm_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.ssm_state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + False, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor，保留每个请求本轮已接受的前缀。""" + + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """非 MTP decode 把单 token 记入当前 history。""" + + self.accept_updates(reqs) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """Prefill 前把已接受 history 折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or SSMRunConfig(32, 4) + bv = config.BV + _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_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, + False, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + self._materialize_accepted( + _upload_indices(req_indices, self.ssm_state.device, torch.int32), + _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """从 acceptance 取出请求和接受长度，供直接持有 history 的测试调用。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), + dtype=self.ssm_state.dtype, + device=self.ssm_state.device, + ) + self._materialize_accepted( + torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py deleted file mode 100644 index ae5a734dba..0000000000 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_runtime.py +++ /dev/null @@ -1,332 +0,0 @@ -"""KDA replay history。只保存 ring buffer、cursor，以及 accepted SSM 的物化和拷贝。""" - -from __future__ import annotations - -from collections.abc import Sequence - -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, -) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import ( - commit_ssm_prefix, -) - - -def upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -class ReplayKDARuntime: - """请求槽上的 KDA replay history。 - - ``state`` 是唯一的 canonical checkpoint；keys、deltas、gates、betas 和 - cursor 只属于请求当前的临时 history，不参与 CPU/PD checkpoint。decode - 先从 checkpoint 加载 history，再把新 token 写入 ring buffer；接受或 - checkpoint 导出时，materialize kernel 将已接受前缀折叠回 ``state``。 - """ - - _run_config: SSMRunConfig | None = None - - @property - def run_config(self) -> SSMRunConfig | None: - """返回物化和 autotune 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" - - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) - - def __init__( - self, - state: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - *, - num_key_heads: int | None = None, - kda: bool = True, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, - ) -> None: - """创建 KDA replay 的 canonical state、双半区 history 和 cursor。 - - ``verify_width`` 决定一次 verify 的最大 token 数，``capacity`` 是 - history ring 的半区长度；``compute_mode`` 选择 projection 的执行方式。 - """ - - assert state.dtype in ( - torch.float32, - torch.bfloat16, - ), "KDA replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - if not kda: - raise ValueError("KDA replay runtime requires the KDA gate") - self.state = state - self.kda = True - 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 - # 保留冗余布局，兼容已有调用方和 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) - # fold 时交替使用两个半区，避免一个 V tile 仍在读取时被其他 CTA 覆盖。 - # cursor 打包保存 phase（2*L）和 count（0..L）。precompute 路径在 BF16 - # 下可以使用更窄的 history buffer；其他形状使用 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), - 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: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM state。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - 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, - 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, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor 以保留每个请求本轮已接受的前缀。 - - MTP 时 ``accepted_state_indices`` 是每个请求最后接受的 draft - 位置；单 token 路径传 ``None``，每个请求只推进一个位置。 - """ - _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: torch.Tensor) -> None: - """把接受的 history 物化回 canonical state，并清空请求 cursor。""" - 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: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 state 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - - layers, slots, hv, k, v = self.state.shape - config = self.run_config or SSMRunConfig(32, 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, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - - def get_accepted_state(self, req_idx: int, accepted_index: int | None = None) -> torch.Tensor: - """返回不修改 active state/history 的单请求 canonical 快照。""" - 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 export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - *, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页，避免中间整批分配。""" - - if not req_indices: - return - shape = (self.state.shape[0], *self.state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.state.dtype for t in targets) - self._materialize_accepted_state( - upload_indices(req_indices, self.state.device, torch.int32), - upload_indices([t.data_ptr() for t in targets], self.state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - - def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 直接把单 token 写入当前 history。""" - self.accept_updates(reqs) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) - - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """兼容直接持有 runtime 的测试：从 acceptance 取出请求和接受长度。""" - - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - - def export_ssm( - self, - req_indices: list[int], - targets: list[torch.Tensor], - *, - flush: bool = False, - ) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先在 GPU 上物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.state.shape[0], *self.state.shape[2:]), - dtype=self.state.dtype, - device=self.state.device, - ) - self._materialize_accepted_state( - torch.tensor([req], dtype=torch.int32, device=self.state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.state[:, req].copy_(snap) - self.clear_history(req) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py index d57dc65038..c4d117826b 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py @@ -33,7 +33,7 @@ def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) for req_idx, slot_id in zip(req_indices, page_slot_ids) ] ) - state.runtime.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + state.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: @@ -44,7 +44,7 @@ def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_i width = state.config.conv_width conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - state.runtime.clear_history(req_idx) + state.clear_history(req_idx) conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) @@ -64,7 +64,7 @@ def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_boun width = state.config.conv_width conv_copies = [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - state.runtime.clear_history(dst_req) + state.clear_history(dst_req) conv_copies.append( ( accepted_conv(remote_dp_bounded_state, src_req), @@ -72,9 +72,7 @@ def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_boun ) ) copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.runtime.export_ssm( - src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices] - ) + remote_dp_bounded_state.export_ssm(src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices]) if state.mtp_state_index is not None: for dst_req in dst_req_indices: state.mtp_state_index[dst_req].fill_(0) @@ -100,7 +98,7 @@ def read_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, lay """读入基准 Conv/SSM，并清空该请求的 history。""" assert layout.state_name == state.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - state.runtime.clear_history(req_idx) + state.clear_history(req_idx) fields = tensor_fields( state.conv_state[:, req_idx, :, : state.config.conv_width], state.ssm_state[:, req_idx], diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 5b41edb4fc..ec43acfedf 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -1,16 +1,26 @@ """KDA replay 的 bounded attention 状态。 -直接实现状态接口，不继承 native KDAState。history 在 replay_kda_runtime， -页布局在 replay_kda_checkpoint，已接受快照在 replay_kda_snapshot。 -一层 decode 重放在 replay_kda_step。 +Conv、canonical SSM 和 replay history 都定义在本类。页布局在 replay_kda_checkpoint， +已接受快照在 replay_kda_snapshot。一层 decode 重放在 replay_kda_step。 """ from __future__ import annotations +from collections.abc import Sequence from functools import lru_cache from typing import NamedTuple import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _accept_updates, + _fold_history, + _materialize, + _prepare_fold, +) +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.ssm_types import SSMRunConfig from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState @@ -22,7 +32,6 @@ read_pages_from_cpu, write_pages_to_cpu, ) -from .replay_kda_runtime import ReplayKDARuntime from .replay_kda_snapshot import ( accepted_conv, copy_runtime_from_dp, @@ -45,8 +54,84 @@ def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: return capacity, compute_mode +def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" + + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: SSMRunConfig | dict[str, int] | None, +) -> None: + """在 owner 上创建 KDA replay 的双半区 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = True + owner.gate_lower_bound = gate_lower_bound + owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.forward_cursors = torch.zeros_like(owner.cursors) + # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + + class ReplayKDAState(BoundedAttentionState): - """KDA replay 运行态。槽位布局和 history 都在本类，不沿用 native 的 MTP 槽。""" + """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + _run_config: SSMRunConfig | None = None + + @property + def state(self) -> torch.Tensor: + """Replay kernel 读取的 canonical SSM，与 ssm_state 是同一张量。""" + + return self.ssm_state + + @property + def run_config(self) -> SSMRunConfig | None: + """返回物化和 autotune 使用的 kernel 配置。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: SSMRunConfig | dict | None) -> None: + """保存 kernel 配置，兼容 autotune 返回的字典。""" + + if config is None or isinstance(config, SSMRunConfig): + self._run_config = config + else: + self._run_config = SSMRunConfig(**config) def __init__( self, @@ -96,17 +181,60 @@ def __init__( self.big_pages_pool = None self.small_pages_pool = None capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) - self.runtime = ReplayKDARuntime( + _bind_history( + self, self.ssm_state, capacity, mtp_step + 1, config.conv_dtype, - num_key_heads=config.num_k_heads, - kda=True, - gate_lower_bound=gate_lower_bound, - compute_mode=compute_mode, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, ) + @classmethod + def from_ssm( + cls, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + *, + num_key_heads: int | None = None, + kda: bool = True, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: SSMRunConfig | dict[str, int] | None = None, + ) -> "ReplayKDAState": + """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" + + if not kda: + raise ValueError("KDA replay requires the KDA gate") + owner = cls.__new__(cls) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "kda" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + @property def ssm_slots_per_req(self) -> int: """每个请求只保留一份 canonical SSM。""" @@ -115,9 +243,9 @@ def ssm_slots_per_req(self) -> int: @property def _ssm_executor(self): - """测试沿用的 history 入口。""" + """History 张量在本对象上。测试沿用这个入口。""" - return self.runtime + return self def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" @@ -188,7 +316,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: for req_idx in req_indices: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() - self.runtime.clear_history(req_idx) + self.clear_history(req_idx) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) @@ -214,15 +342,11 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to """提交已接受前缀，并推进 replay cursor。""" if self.mtp_state_index is None: return - self.runtime.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """Prefill 前把已接受 history 折回 canonical state。""" - self.runtime.merge_accepted_updates(reqs) + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(accepted_conv(self, req_idx), self.runtime.get_accepted_state(req_idx)) + return _AcceptedTensors(accepted_conv(self, req_idx), self.accepted_ssm(req_idx)) def copy_runtime_state_from_dp( self, @@ -251,3 +375,201 @@ def read_from_pd_page( layout, ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.ssm_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, + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.ssm_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.ssm_state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + True, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor，保留每个请求本轮已接受的前缀。""" + + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """非 MTP decode 把单 token 记入当前 history。""" + + self.accept_updates(reqs) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """Prefill 前把已接受 history 折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or SSMRunConfig(32, 4) + bv = config.BV + _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_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, + True, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + *, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + self._materialize_accepted( + _upload_indices(req_indices, self.ssm_state.device, torch.int32), + _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: + """从 acceptance 取出请求和接受长度，供直接持有 history 的测试调用。""" + + self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), + dtype=self.ssm_state.dtype, + device=self.ssm_state.device, + ) + self._materialize_accepted( + torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index e5ebe6d83a..12fb0f55ee 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -247,7 +247,7 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) self._init_indices(draft_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - begin_decode(state.runtime, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state @@ -321,7 +321,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat backend.head_v_dim, ) output, _ = replay( - state.runtime, + state, state.storage_layer(layer_num), query, key, @@ -350,7 +350,7 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step ) query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) output, _ = replay( - state.runtime, + state, state.storage_layer(layer_num), query, key, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py index db887cb923..89af883103 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py @@ -1,6 +1,6 @@ """GDN replay 的一层 decode。 -History 张量和 cursor 留在 ReplayGDNRuntime。这里只选择 kernel 配置并发起重放。 +History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置并发起重放。 """ from __future__ import annotations @@ -56,7 +56,7 @@ def replay( if history.kda: raise ValueError("GDN replay step does not use the KDA gate") config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - hv, kd, vd = history.state.shape[-3:] + hv, kd, vd = history.ssm_state.shape[-3:] assert history.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) @@ -72,7 +72,7 @@ def replay( b, a_log, bias, - history.state[linear_layer_idx], + history.ssm_state[linear_layer_idx], history.keys[linear_layer_idx], history.deltas[linear_layer_idx], history.gates[linear_layer_idx], diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index d8325fcf20..999c51a8cd 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -159,7 +159,7 @@ def init_state(self): state = _single_replay_state(self) self._init_indices(self.backend.mtp_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - begin_decode(state.runtime, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state @@ -230,7 +230,7 @@ def decode_att( raw_gate = params["raw_gate"].view(*shape, backend.tp_hidden_size) raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) output, _ = replay( - state.runtime, + state, state.storage_layer(layer_num), q, k, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_step.py b/lightllm/common/basemodel/attention/linear/replay_kda_step.py index 5fcfa29d45..78aaa8ce43 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_step.py @@ -1,6 +1,6 @@ """KDA replay 的一层 decode。 -History 张量和 cursor 留在 ReplayKDARuntime。这里只选择 kernel 配置并发起重放。 +History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置并发起重放。 """ from __future__ import annotations @@ -56,7 +56,7 @@ def replay( if not history.kda: raise ValueError("KDA replay step requires the KDA gate") config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - hv, kd, vd = history.state.shape[-3:] + hv, kd, vd = history.ssm_state.shape[-3:] assert history.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) @@ -72,7 +72,7 @@ def replay( b, a_log, bias, - history.state[linear_layer_idx], + history.ssm_state[linear_layer_idx], history.keys[linear_layer_idx], history.deltas[linear_layer_idx], history.gates[linear_layer_idx], diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 2556e1681b..fe9861bec1 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -89,7 +89,15 @@ def rebuild_inputs( ) _, _, 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) - scratch = type(ssm_backend)( + if ssm_backend.kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + + history_cls = ReplayKDAState + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + history_cls = ReplayGDNState + scratch = history_cls.from_ssm( state, ssm_backend.capacity, ssm_backend.verify_width, diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 14ea97b744..fd866a87e4 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -20,12 +20,12 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime + return ReplayKDAState.from_ssm(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - return ReplayGDNRuntime(*args, **kwargs) + return ReplayGDNState.from_ssm(*args, **kwargs) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer 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 dc640c029d..0a8e0e5c40 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -80,8 +80,8 @@ def test_hybrid_layer_mapping_and_mtp_states( ) assert bounded_state.ssm_slots_per_req == (3 if mode == "native" else 1) if mode == "replay": - assert bounded_state.runtime.state is bounded_state.ssm_state - assert bounded_state.runtime.capacity == 8 + assert bounded_state.state is bounded_state.ssm_state + assert bounded_state.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) manager.register_state(model_role="main", state=bounded_state) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index f71243a216..46b450bd21 100644 --- 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 @@ -8,17 +8,21 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime + return ReplayKDAState.from_ssm(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - return ReplayGDNRuntime(*args, **kwargs) + return ReplayGDNState.from_ssm(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) from types import SimpleNamespace from lightllm.common.attention_state.commit import StateAcceptance -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer @@ -36,7 +40,7 @@ def make(n, width): 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 = ReplayGDNRuntime.__new__(ReplayGDNRuntime) + ssm_backend = ReplayGDNState.__new__(ReplayGDNState) ssm_backend.cursors = tensor([7] * (hold + 1)) ssm_backend.hold_req_id = hold ssm_backend.capacity = max(4, triton.next_power_of_2(width)) @@ -225,7 +229,7 @@ def cycle(i): 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)) + assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) if iteration % 11 == 10: for ssm_backend in ssm_backends: ssm_backend.merge_accepted_updates(reqs[:3]) 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 index dccd92afba..f2b8c9c9a7 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -5,12 +5,16 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime + return ReplayKDAState.from_ssm(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - return ReplayGDNRuntime(*args, **kwargs) + return ReplayGDNState.from_ssm(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer @@ -89,7 +93,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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) + torch.testing.assert_close(ssm_snapshot(ssm_backend, 0)[0].float(), expected_state, atol=128, rtol=0.01) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -107,7 +111,7 @@ def test_flush_export_updates_checkpoint_and_clears_history(dtype): backend.gates.fill_(-0.1) backend.betas.fill_(0.5) backend.cursors[0] = 2 - expected = backend.get_accepted_state(0).clone() + expected = ssm_snapshot(backend, 0).clone() target = torch.empty_like(expected) backend.export_accepted_state([0], [target], flush=True) @@ -474,7 +478,7 @@ def step(): def test_snapshot_without_history_copies_active_state(): state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) ssm_backend = ReplaySSMExecutor(state, 4, 3) - snapshot = ssm_backend.get_accepted_state(1) + snapshot = ssm_snapshot(ssm_backend, 1) torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) assert ssm_backend.cursors[1].item() == 0 @@ -529,7 +533,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg backend.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) ) expected += key[..., None] * delta[..., None, :] - materialized = backend.get_accepted_state(1) + materialized = ssm_snapshot(backend, 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) 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 index 77436ca1ff..07cca89b23 100644 --- 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 @@ -6,12 +6,12 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime + return ReplayKDAState.from_ssm(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - return ReplayGDNRuntime(*args, **kwargs) + return ReplayGDNState.from_ssm(*args, **kwargs) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer 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 index 2ee7477f54..7648245bc0 100644 --- 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 @@ -8,12 +8,12 @@ def ReplaySSMExecutor(*args, **kwargs): kda = kwargs.get("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_runtime import ReplayKDARuntime + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - return ReplayKDARuntime(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime + return ReplayKDAState.from_ssm(*args, **kwargs) + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - return ReplayGDNRuntime(*args, **kwargs) + return ReplayGDNState.from_ssm(*args, **kwargs) @pytest.mark.parametrize("width", [1, 5]) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index 2c8163fd2d..b7c060f711 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -20,9 +20,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): state = getattr(resources, state_name) assert state.big_pages_pool is state.small_pages_pool is None assert resources.full_att_allocator is None - runtime = getattr(state, "runtime", None) - if runtime is not None: - assert runtime.state is state.ssm_state + if mode == "replay": + assert state.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) @@ -86,7 +85,9 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): 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 getattr(copy, "runtime", None) is getattr(state, "runtime", None) + if mode == "replay": + assert copy.cursors is state.cursors + assert copy.keys is state.keys assert copy.mtp_state_index is state.mtp_state_index assert all(vars(state)[name] is value for name, value in before.items()) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 88e9504964..5120da1fee 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -138,14 +138,16 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty for req in requests ] before = state.ssm_state.clone() - cursors = state.runtime.cursors.clone() if mode == "replay" else None + history = getattr(state, "runtime", state) + cursors = history.cursors.clone() if mode == "replay" else None def no_snapshot(*args): raise AssertionError("batch export must write directly to the destination") if mode == "replay": monkeypatch.setattr(state, "get_accepted_state", no_snapshot) - monkeypatch.setattr(state.runtime, "get_accepted_state", no_snapshot) + if history is not state: + monkeypatch.setattr(history, "get_accepted_state", no_snapshot) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 7e0e6fb870..c4a62a836a 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -8,7 +8,6 @@ from lightllm.common.attention_state.bounded_att import KDAState from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_runtime import ReplayGDNRuntime from lightllm.common.basemodel.attention.linear.replay_gdn_step import run_config as gdn_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState @@ -89,11 +88,11 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): def test_kda_backend_receives_gate_bound_at_construction(linear_config): owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") - assert owner.runtime.gate_lower_bound == -2.5 + assert owner.gate_lower_bound == -2.5 def test_replay_backend_publishes_detached_run_config(): - ssm_backend = object.__new__(ReplayGDNRuntime) + ssm_backend = object.__new__(ReplayGDNState) tuning = {"BV": 32, "num_warps": 2} ssm_backend.run_config = tuning ssm_backend._config_is_fixed = True @@ -112,7 +111,8 @@ def test_replay_backend_publishes_detached_run_config(): def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") - assert owner.runtime.compute_mode == expected + history = getattr(owner, "runtime", owner) + assert history.compute_mode == expected @pytest.mark.parametrize("cls", [GDNState, KDAState]) @@ -131,14 +131,13 @@ def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): import sys owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + history = getattr(owner, "runtime", owner) calls = [] - monkeypatch.setattr( - sys.modules[owner.runtime.__class__.__module__], "commit_ssm_prefix", lambda *args: calls.append(args) - ) - monkeypatch.setattr(owner.runtime, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) + monkeypatch.setattr(history, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3])) assert len(calls) == 1 - assert calls[0][3] is owner.runtime.cursors + assert calls[0][3] is history.cursors assert calls[0][0] is owner.mtp_state_index @@ -148,9 +147,10 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): pytest.skip("CUDA required") owner, _, _ = make_state(mode="replay", step=3, device=device) owner.mtp_state_index.fill_(2) - owner.runtime.cursors.fill_(8) + history = getattr(owner, "runtime", owner) + history.cursors.fill_(8) reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] owner.commit_bounded_att_state(reqs, lengths) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] - assert owner.runtime.cursors.tolist() == [11, 8, 9, 8] + assert history.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index f11647a5e9..b882a27c88 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -152,7 +152,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear._ssm_executor.gates.fill_(-0.5) linear._ssm_executor.betas.zero_() linear._ssm_executor.cursors[0] = 1 - ssm = linear._ssm_executor.get_accepted_state(0).clone() + ssm = linear.accepted_ssm(0).clone() linear._ssm_executor.cursors[2] = 3 if mtp_step: linear.mtp_state_index[2] = mtp_step diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 24ccc83819..704849ecba 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -296,7 +296,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod selected._ssm_executor.accept_updates(ids, accepted) for req, length in zip(requests, lengths): torch.testing.assert_close( - selected._ssm_executor.get_accepted_state(req)[0], + selected.accepted_ssm(req)[0], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 7a3884a58d..c76091af1e 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -173,7 +173,7 @@ def test_pd_kv_and_runtime_state_roundtrip( backend.gates.fill_(-0.5) backend.betas.zero_() backend.cursors[src_req] = 1 - ssm[:, heads].copy_(backend.get_accepted_state(src_req)) + ssm[:, heads].copy_(mem.linear.accepted_ssm(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) From 93d7249d90b34f3c3fa2ddf6e4fc95e7c454973f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 05:46:07 +0000 Subject: [PATCH 11/54] refactor: fold replay history when decode accepts Prefill writes the canonical SSM directly. Replay keeps its kernel config local and no longer shares a slot stride or an executor alias. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 70 +++++++++---------- .../bounded_att/kda/replay_kda_state.py | 70 +++++++++---------- .../basemodel/attention/linear/replay_gdn.py | 7 +- .../attention/linear/replay_gdn_step.py | 13 ++-- .../basemodel/attention/linear/replay_kda.py | 5 +- .../attention/linear/replay_kda_step.py | 13 ++-- .../triton_kernel/linear_att/ssm_autotune.py | 57 ++++++++------- lightllm/common/ssm_types.py | 14 ---- test/benchmark/kernels/benchmark_kda_ssm.py | 3 +- .../test_hybrid_layer_mapping.py | 4 +- .../linear_att/test_acceptance_fusion.py | 10 +-- .../linear_att/test_replayssm.py | 35 +++++----- .../linear_att/test_replayssm_execution.py | 29 +++++--- .../linear_att/test_ssm_autotune_memory.py | 3 +- .../test_hybrid_req_manager_pd_sharing.py | 21 +++--- unit_tests/common/test_linear_state.py | 37 +++++----- .../common/test_linear_state_contract.py | 30 +------- unit_tests/models/glm5_next/test_cache.py | 22 +++--- unit_tests/models/glm5_next/test_mtp.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 11 +-- 20 files changed, 218 insertions(+), 238 deletions(-) delete mode 100644 lightllm/common/ssm_types.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 7abfbada28..ab4b2c9a55 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -20,8 +20,6 @@ _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from lightllm.common.ssm_types import SSMRunConfig - from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig @@ -69,7 +67,7 @@ def _bind_history( num_key_heads: int | None, gate_lower_bound: float, compute_mode: str, - run_config: SSMRunConfig | dict[str, int] | None, + run_config: dict[str, int] | None, ) -> None: """在 owner 上创建 GDN replay 的双半区 history 和 cursor。""" @@ -110,28 +108,26 @@ def _bind_history( class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" - _run_config: SSMRunConfig | None = None + _run_config: dict[str, int] | None = None @property - def state(self) -> torch.Tensor: - """Replay kernel 读取的 canonical SSM，与 ssm_state 是同一张量。""" - - return self.ssm_state - - @property - def run_config(self) -> SSMRunConfig | None: + def run_config(self) -> dict[str, int] | None: """返回物化和 autotune 使用的 kernel 配置。""" return self._run_config @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } def __init__( self, @@ -205,7 +201,7 @@ def from_ssm( kda: bool = False, gate_lower_bound: float = -5.0, compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, + run_config: dict[str, int] | None = None, ) -> "ReplayGDNState": """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" @@ -235,18 +231,6 @@ def from_ssm( ) return owner - @property - def ssm_slots_per_req(self) -> int: - """每个请求只保留一份 canonical SSM。""" - - return 1 - - @property - def _ssm_executor(self): - """History 张量在本对象上。测试沿用这个入口。""" - - return self - def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" return ( @@ -339,10 +323,12 @@ def restore_bounded_att_runtime_from_pages( restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受前缀，并推进 replay cursor。""" + """提交已接受前缀，折进 canonical SSM，并清空 history。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + if self.ssm_state.is_cuda: + self.merge_accepted_updates(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -376,6 +362,18 @@ def read_from_pd_page( ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, *, fold: bool) -> None: + """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + + if self.config is None or not self.config.logical_layer_ids: + return + if logical_layer_id != self.config.logical_layer_ids[-1]: + return + if fold: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -455,7 +453,7 @@ def accept_single_token(self, reqs: torch.Tensor) -> None: self.accept_updates(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """Prefill 前把已接受 history 折回 canonical SSM，并清空 cursor。""" + """把已接受 history 折回 canonical SSM，并清空 cursor。""" if reqs.numel() == 0: return @@ -474,8 +472,8 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape - config = self.run_config or SSMRunConfig(32, 4) - bv = config.BV + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, @@ -498,8 +496,8 @@ def _materialize_accepted( False, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, + num_warps=config["num_warps"], + num_stages=config["num_stages"], ) def accepted_ssm(self, req_idx: int) -> torch.Tensor: diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index ec43acfedf..56d573eef6 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -20,8 +20,6 @@ _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from lightllm.common.ssm_types import SSMRunConfig - from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig @@ -69,7 +67,7 @@ def _bind_history( num_key_heads: int | None, gate_lower_bound: float, compute_mode: str, - run_config: SSMRunConfig | dict[str, int] | None, + run_config: dict[str, int] | None, ) -> None: """在 owner 上创建 KDA replay 的双半区 history 和 cursor。""" @@ -110,28 +108,26 @@ def _bind_history( class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" - _run_config: SSMRunConfig | None = None + _run_config: dict[str, int] | None = None @property - def state(self) -> torch.Tensor: - """Replay kernel 读取的 canonical SSM，与 ssm_state 是同一张量。""" - - return self.ssm_state - - @property - def run_config(self) -> SSMRunConfig | None: + def run_config(self) -> dict[str, int] | None: """返回物化和 autotune 使用的 kernel 配置。""" return self._run_config @run_config.setter - def run_config(self, config: SSMRunConfig | dict | None) -> None: - """保存 kernel 配置，兼容 autotune 返回的字典。""" + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" - if config is None or isinstance(config, SSMRunConfig): - self._run_config = config - else: - self._run_config = SSMRunConfig(**config) + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } def __init__( self, @@ -205,7 +201,7 @@ def from_ssm( kda: bool = True, gate_lower_bound: float = -5.0, compute_mode: str = "inline", - run_config: SSMRunConfig | dict[str, int] | None = None, + run_config: dict[str, int] | None = None, ) -> "ReplayKDAState": """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" @@ -235,18 +231,6 @@ def from_ssm( ) return owner - @property - def ssm_slots_per_req(self) -> int: - """每个请求只保留一份 canonical SSM。""" - - return 1 - - @property - def _ssm_executor(self): - """History 张量在本对象上。测试沿用这个入口。""" - - return self - def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" return ( @@ -339,10 +323,12 @@ def restore_bounded_att_runtime_from_pages( restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受前缀，并推进 replay cursor。""" + """提交已接受前缀，折进 canonical SSM，并清空 history。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + if self.ssm_state.is_cuda: + self.merge_accepted_updates(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -376,6 +362,18 @@ def read_from_pd_page( ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, *, fold: bool) -> None: + """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + + if self.config is None or not self.config.logical_layer_ids: + return + if logical_layer_id != self.config.logical_layer_ids[-1]: + return + if fold: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -455,7 +453,7 @@ def accept_single_token(self, reqs: torch.Tensor) -> None: self.accept_updates(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """Prefill 前把已接受 history 折回 canonical SSM，并清空 cursor。""" + """把已接受 history 折回 canonical SSM，并清空 cursor。""" if reqs.numel() == 0: return @@ -474,8 +472,8 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape - config = self.run_config or SSMRunConfig(32, 4) - bv = config.BV + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, @@ -498,8 +496,8 @@ def _materialize_accepted( True, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, - num_warps=config.num_warps, - num_stages=config.num_stages, + num_warps=config["num_warps"], + num_stages=config["num_stages"], ) def accepted_ssm(self, req_idx: int) -> torch.Tensor: diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 12fb0f55ee..6cd5b38c17 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -181,8 +181,8 @@ class ReplayGDNPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): + # Prefill 直接覆盖 canonical SSM。进入前 history 已经在 decode 接受时折完。 reqs = self.infer_state.b_req_idx - _single_replay_state(self).merge_accepted_updates(reqs) self.b_conv_buffer_idx = reqs self.b_ssm_buffer_idx = reqs @@ -204,7 +204,9 @@ def prefill_att( if backend.mtp_step > 0: conv_states = conv_states[:, :, : -backend.mtp_step] mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - return self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight), z + core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) + state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=False) + return core_attn_out, z def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): backend = self.backend @@ -333,6 +335,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) + state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=True) return output, z def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py index 89af883103..33f09c831d 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_step.py @@ -10,7 +10,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config -from lightllm.common.ssm_types import SSMRunConfig def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: @@ -31,7 +30,7 @@ def run_config( a_log: torch.Tensor, bias: torch.Tensor, cu_seqlens: torch.Tensor | None = None, -) -> SSMRunConfig: +) -> dict[str, int]: """按当前 kernel 形状选择或返回固定的 GDN replay 调优配置。""" return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) @@ -49,8 +48,8 @@ def replay( bias: torch.Tensor, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, -) -> tuple[torch.Tensor, SSMRunConfig]: + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: """执行一层 GDN replay decode，并返回输出与本轮 run config。""" if history.kda: @@ -60,7 +59,7 @@ def replay( assert history.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 + bv = config["BV"] precompute_projections = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) @@ -104,7 +103,7 @@ def replay( history.gate_lower_bound, BQK=max(16, 2 * bw), BR=triton.next_power_of_2(max(16, history.capacity + bw)), - num_warps=config.num_warps, - num_stages=config.num_stages, + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), ) return out, config diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 999c51a8cd..d3539311ba 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -92,8 +92,8 @@ class ReplayKDAPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): + # Prefill 直接覆盖 canonical SSM。进入前 history 已经在 decode 接受时折完。 reqs = self.infer_state.b_req_idx - _single_replay_state(self).merge_accepted_updates(reqs) self.b_conv_buffer_idx = reqs self.b_ssm_buffer_idx = reqs @@ -145,6 +145,7 @@ def prefill_att( lower_bound=backend.lower_bound, ) ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=False) return output @@ -242,4 +243,6 @@ def decode_att( self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) + if self.b1_mtp_cu_q_seq_len is None: + state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=True) return output diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_step.py b/lightllm/common/basemodel/attention/linear/replay_kda_step.py index 78aaa8ce43..770f74ee3a 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_step.py @@ -10,7 +10,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config -from lightllm.common.ssm_types import SSMRunConfig def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: @@ -31,7 +30,7 @@ def run_config( a_log: torch.Tensor, bias: torch.Tensor, cu_seqlens: torch.Tensor | None = None, -) -> SSMRunConfig: +) -> dict[str, int]: """按当前 kernel 形状选择或返回固定的 KDA replay 调优配置。""" return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) @@ -49,8 +48,8 @@ def replay( bias: torch.Tensor, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None, - run_config: SSMRunConfig | None = None, -) -> tuple[torch.Tensor, SSMRunConfig]: + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: """执行一层 KDA replay decode，并返回输出与本轮 run config。""" if not history.kda: @@ -60,7 +59,7 @@ def replay( assert history.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 + bv = config["BV"] precompute_projections = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) @@ -104,7 +103,7 @@ def replay( history.gate_lower_bound, BQK=max(16, 2 * bw), BR=triton.next_power_of_2(max(16, history.capacity + bw)), - num_warps=config.num_warps, - num_stages=config.num_stages, + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), ) return out, config diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index fe9861bec1..84b96964ff 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -9,10 +9,11 @@ from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune from lightllm.utils.log_utils import init_logger -from lightllm.common.ssm_types import SSMRunConfig logger = init_logger(__name__) +_DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} + def get_configs() -> list[dict[str, int]]: """返回 ReplaySSM update cycle 的候选 Triton 配置。""" @@ -26,13 +27,13 @@ def get_configs() -> list[dict[str, int]]: def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens) -> dict[str, object]: """构造区分状态布局、数据类型和投影模式的静态调优键。""" axis = 1 if cu_seqlens is not None else 0 - _, _, hv, kd, vd = ssm_backend.state.shape + _, _, hv, kd, vd = ssm_backend.ssm_state.shape # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 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)), + "dtype": str((q.dtype, ssm_backend.ssm_state.dtype)), "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), "history": ssm_backend.capacity, "gate": int(ssm_backend.kda), @@ -58,7 +59,7 @@ def run_key(ssm_backend, q, cu_seqlens) -> tuple[int, int]: def scratch_size_bytes(ssm_backend, batch, sequences, inputs) -> int: """估算一次 SSM autotune scratch 所需的字节数。""" # State/history 的首轴是 layer，cursor 在各 layer 之间共享。 - layers, slots = ssm_backend.state.shape[:2] + layers, slots = ssm_backend.ssm_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() @@ -87,17 +88,9 @@ def rebuild_inputs( 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 ssm_backend.kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - - history_cls = ReplayKDAState - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - history_cls = ReplayGDNState - scratch = history_cls.from_ssm( + _, _, hv, kd, vd = ssm_backend.ssm_state.shape + state = torch.empty((1, batch + 1, hv, kd, vd), device=q.device, dtype=ssm_backend.ssm_state.dtype).normal_(0, 0.02) + scratch = type(ssm_backend).from_ssm( state, ssm_backend.capacity, ssm_backend.verify_width, @@ -138,12 +131,17 @@ def rebuild_inputs( def replay_layer(history, *args, **kwargs): - """按 history 所属算法调用对应的 replay decode。""" + """按 history 的具体类型调用对应的 replay decode。""" + + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - if history.kda: + if isinstance(history, ReplayKDAState): from lightllm.common.basemodel.attention.linear.replay_kda_step import replay - else: + elif isinstance(history, ReplayGDNState): from lightllm.common.basemodel.attention.linear.replay_gdn_step import replay + else: + raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(history).__name__}") return replay(history, *args, **kwargs) @@ -159,7 +157,7 @@ def replay_layer(history, *args, **kwargs): ) def select_config( ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None -) -> dict[str, int] | SSMRunConfig | None: +) -> dict[str, int] | None: """搜索 update cycle 布局；计时只使用重建的临时输入。""" if workload is None: return run_config @@ -176,18 +174,27 @@ def select_config( return run_config -def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> SSMRunConfig: +def _replay_config(config: dict | None) -> dict[str, int]: + if not config: + return dict(_DEFAULT_REPLAY_CONFIG) + return { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + + +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: """返回当前 batch 的固定或已调优 SSM kernel 配置。""" - default = ssm_backend.run_config or SSMRunConfig(32, 1) - if ssm_backend._config_is_fixed: - return default + if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: + return ssm_backend.run_config + default = _replay_config(ssm_backend.run_config) 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 = SSMRunConfig(**config) if config else default + config = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 ssm_backend.batch_configs[key] = config logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") diff --git a/lightllm/common/ssm_types.py b/lightllm/common/ssm_types.py deleted file mode 100644 index a0d4da1389..0000000000 --- a/lightllm/common/ssm_types.py +++ /dev/null @@ -1,14 +0,0 @@ -from dataclasses import dataclass - - -@dataclass(frozen=True, slots=True) -class SSMRunConfig: - BV: int - num_warps: int - num_stages: int = 3 - - def __post_init__(self) -> None: - """校验 Triton update cycle 配置的字段类型。""" - 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/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index fd866a87e4..ea9020dd18 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -4,7 +4,6 @@ HTTP scheduling. Use a matched service workload for end-to-end conclusions. """ -from dataclasses import asdict import argparse import gc import json @@ -149,7 +148,7 @@ def main(): state_history_bytes=size, ms_per_round=statistics.median(samples), samples_ms=samples, - config=asdict(ssm_backend.run_config) if ssm_backend else None, + config=dict(ssm_backend.run_config) if ssm_backend and ssm_backend.run_config else None, candidates=candidates, ) rows.append(row) 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 0a8e0e5c40..2b98f3da97 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -78,9 +78,9 @@ 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) + if mode == "native": + assert bounded_state.ssm_slots_per_req == 3 if mode == "replay": - assert bounded_state.state is bounded_state.ssm_state assert bounded_state.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) 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 index 46b450bd21..20b1fa6a12 100644 --- 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 @@ -225,7 +225,7 @@ def cycle(i): 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].ssm_state, ssm_backends[1].ssm_state) assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) if iteration % 7 == 0: @@ -233,13 +233,13 @@ def cycle(i): 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) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) if iteration == 19: for manager in managers: manager._ssm_executor.clear_history(1) - manager._ssm_executor.state[:, 1].zero_() + manager._ssm_executor.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]]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index f2b8c9c9a7..362b7024a2 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -351,7 +351,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d 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) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) def test_bf16_decode_fold_uses_rounded_checkpoint(): @@ -380,7 +380,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): ssm_backend.accept_updates(reqs) outputs.append(replay_layer(ssm_backend, 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) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -512,28 +512,27 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg state.mtp_state_index[1] = 1 state.conv_state.normal_() state.ssm_state.zero_() - backend = state._ssm_executor - backend.raw_keys.normal_() - backend.raw_values.normal_() - backend.gates.uniform_(-0.1, -0.01) - backend.betas.uniform_(0.1, 0.9) + state.raw_keys.normal_() + state.raw_values.normal_() + state.gates.uniform_(-0.1, -0.01) + state.betas.uniform_(0.1, 0.9) active_before = state.ssm_state[:, 1].clone() conv_expected = state.conv_state[:, 1, :, 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 + state.cursors[1] = count expected = active_before.float().clone() for i in range(count): - key = backend.raw_keys[:, 1, :, i].float() + key = state.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 *= state.gates[:, 1, :, i].exp()[..., None, None] + delta = state.betas[:, 1, :, i, None] * ( + state.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) ) expected += key[..., None] * delta[..., None, :] - materialized = ssm_snapshot(backend, 1) + materialized = ssm_snapshot(state, 1) torch.testing.assert_close(materialized, expected.to(dtype), rtol=state_rtol, atol=2e-5) state.save_bounded_att_runtime_to_pages([1], [8], kind, [0]) state.write_to_pd_page(1, block, 0, 8, layout) @@ -541,15 +540,15 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg 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 + assert state.cursors[1].item() == count for restore in ( lambda: state.restore_bounded_att_runtime_from_pages([2], [8], kind, [0]), lambda: state.read_from_pd_page(block, 2, 0, 8, layout), ): - backend.cursors[2] = 3 + state.cursors[2] = 3 state.mtp_state_index[2] = 2 restore() - assert backend.cursors[2].item() == state.mtp_state_index[2].item() == 0 + assert state.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]) @@ -669,5 +668,5 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) backend.commit(StateAcceptance(reqs[:live], lengths), indices, "main") backend.merge_accepted_updates(reqs[:live]) - torch.testing.assert_close(backend.state[0, :live], expected, atol=3e-4, rtol=3e-4) - assert torch.equal(backend.state[0, live], initial[0, live]) + torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) + assert torch.equal(backend.ssm_state[0, live], initial[0, live]) 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 index 07cca89b23..3c2a81ee92 100644 --- 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 @@ -70,7 +70,7 @@ def cycle(backend, ids): with torch.cuda.graph(graph): out = cycle(captured, ids) graphs[bucket] = graph, ids, out - captured.state.copy_(initial) + captured.ssm_state.copy_(initial) captured.cursors.zero_() counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) for count in counts: @@ -84,12 +84,12 @@ def cycle(backend, ids): assert len(launches) == calls torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) assert torch.equal(captured.cursors, eager.cursors) - torch.testing.assert_close(captured.state, eager.state, atol=3e-5, rtol=3e-5) + torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) assert len(hashes["_prepare_fold"]) == len({triton.next_power_of_2(n) for n in counts}) for name in kernels[1:]: assert len(hashes[name]) == 1, (name, hashes[name]) assert captured.cursors[-1] == 0 - assert torch.equal(captured.state[:, -1], initial[:, -1]) + assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -174,7 +174,7 @@ def forward_accept(index, cu, reqs): 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.ssm_state.copy_(initial) ssm_backend.cursors.zero_() manager.mtp_state_index.zero_() g = torch.cuda.CUDAGraph() @@ -186,7 +186,7 @@ def forward_accept(index, cu, reqs): captured.append((cu, reqs, pos)) outputs.append(out) for ssm_backend in ssm_backends: - ssm_backend.state.copy_(initial) + ssm_backend.ssm_state.copy_(initial) ssm_backend.cursors.zero_() manager.mtp_state_index.zero_() folds = 0 @@ -210,7 +210,10 @@ def forward_accept(index, cu, reqs): 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 + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, ) table = manager.mtp_state_index.cpu().tolist() for req, count in zip(ids, counts): @@ -262,7 +265,7 @@ def warm(ssm_backend): graphs, positions, outputs = [], [], [] for early, ssm_backend in enumerate(ssm_backends): warm(ssm_backend) - ssm_backend.state.copy_(initial) + ssm_backend.ssm_state.copy_(initial) ssm_backend.cursors.zero_() if stagger: reqs.copy_(reqs.new_tensor([0, 2, 2])) @@ -303,7 +306,10 @@ def warm(ssm_backend): 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 + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_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 @@ -398,7 +404,7 @@ def step(): out = step() graphs.append(graph) outputs.append(out) - ssm_backend.state.copy_(initial) + ssm_backend.ssm_state.copy_(initial) ssm_backend.cursors.zero_() for iteration in range(32): accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 @@ -410,5 +416,8 @@ def step(): 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 + ssm_backends[0].ssm_state.float(), + ssm_backends[1].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, ) 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 index 7648245bc0..2627540fb2 100644 --- 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 @@ -1,6 +1,5 @@ import pytest import torch -from lightllm.common.ssm_types import SSMRunConfig from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning @@ -69,7 +68,7 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): config = tuning.resolve_run_config(**inputs) assert tuning.resolve_run_config(**inputs) is config - assert config == SSMRunConfig(32, 1) + assert config == {"BV": 32, "num_warps": 1, "num_stages": 3} 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) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index b7c060f711..807aac36c2 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -20,8 +20,6 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): 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 mode == "replay": - assert state.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) @@ -36,8 +34,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): assert state.mtp_state_index[2] == 0 assert state.mtp_state_index[1] == 1 if mode == "replay": - assert state._ssm_executor.cursors[2] == 0 - assert state._ssm_executor.cursors[1] == 1 + assert state.cursors[2] == 0 + assert state.cursors[1] == 1 torch.cuda.synchronize() @@ -106,11 +104,11 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( state.ssm_state[:, slot].fill_(slot + 1) state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 if mode == "replay": - state._ssm_executor.raw_keys.zero_() - state._ssm_executor.raw_values.zero_() - state._ssm_executor.gates.fill_(-0.5) - state._ssm_executor.betas.zero_() - state._ssm_executor.cursors[1], state._ssm_executor.cursors[2] = 1, 3 + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[1], state.cursors[2] = 1, 3 if hasattr(state, "get_accepted_state"): expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] else: @@ -120,7 +118,8 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( state.ssm_state[:, state.ssm_slots_per_req].cpu().tolist(), ] before = vars(state).copy() - source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, state.ssm_slots_per_req].clone() + source_slot = 1 if mode == "replay" else state.ssm_slots_per_req + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() snapshots = [] for pool in (state.big_pages_pool, state.small_pages_pool): for tensor in (pool.conv_state, pool.ssm_state): @@ -139,7 +138,7 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( 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) + torch.testing.assert_close(state.ssm_state[:, source_slot], source_ssm, rtol=0, atol=0) for tensor, pointer, contents in snapshots: assert tensor.data_ptr() == pointer torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 5120da1fee..b929b40afd 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -56,11 +56,11 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() expected_ssm = state.ssm_state[:, src].clone() if mode == "replay": - state._ssm_executor.raw_keys.zero_() - state._ssm_executor.raw_values.zero_() - state._ssm_executor.gates.fill_(-0.5) - state._ssm_executor.betas.zero_() - state._ssm_executor.cursors[src] = 1 + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.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") @@ -73,7 +73,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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_executor.cursors[src].item() == 1 + assert state.cursors[src].item() == 1 for restore in ( lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), @@ -83,7 +83,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if state.mtp_state_index is not None: state.mtp_state_index[dst] = step if mode == "replay": - state._ssm_executor.cursors[dst] = 3 + state.cursors[dst] = 3 restore() if mode == "replay": actual = state.get_accepted_state(dst) @@ -96,11 +96,14 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if state.mtp_state_index is not None: assert state.mtp_state_index[dst].item() == 0 if mode == "replay": - assert state._ssm_executor.cursors[dst].item() == 0 + assert state.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() + if mode == "replay": + assert not state.ssm_state[:, dst].any() + else: + 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): @@ -121,11 +124,11 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty state.ssm_state[:, slot].fill_(slot + 1) state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) if mode == "replay": - state._ssm_executor.raw_keys.zero_() - state._ssm_executor.raw_values.zero_() - state._ssm_executor.gates.fill_(-0.5) - state._ssm_executor.betas.zero_() - state._ssm_executor.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) requests = [2, 0, 1] if mode == "replay": expected = [(item.conv.clone(), item.ssm.clone()) for item in map(state.get_accepted_state, requests)] @@ -159,7 +162,7 @@ def no_snapshot(*args): torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) assert torch.equal(state.ssm_state, before) if cursors is not None: - assert torch.equal(state._ssm_executor.cursors, cursors) + assert torch.equal(state.cursors, cursors) target, _, _ = make_state(cls, mode, 2, dtype) if target.mtp_state_index is not None: target.mtp_state_index.fill_(2) @@ -175,7 +178,7 @@ def no_snapshot(*args): torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) assert torch.equal(state.ssm_state, before) if cursors is not None: - assert torch.equal(state._ssm_executor.cursors, cursors) + assert torch.equal(state.cursors, cursors) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index c4a62a836a..ba2ec508db 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -1,10 +1,6 @@ -import pickle -from dataclasses import FrozenInstanceError - import pytest import torch -from lightllm.common.ssm_types import SSMRunConfig 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 @@ -15,24 +11,6 @@ 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( @@ -98,12 +76,8 @@ def test_replay_backend_publishes_detached_run_config(): ssm_backend._config_is_fixed = True ssm_backend.compute_mode = "inline" config = gdn_run_config(ssm_backend, *([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 + assert config == {"BV": 32, "num_warps": 2, "num_stages": 3} @pytest.mark.parametrize("cls", [GDNState, KDAState]) @@ -151,6 +125,6 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): history.cursors.fill_(8) reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] - owner.commit_bounded_att_state(reqs, lengths) + owner.commit_accepted(owner.mtp_state_index, reqs, lengths) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] assert history.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index b882a27c88..36431fb3cd 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -147,13 +147,13 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() if mode == "replay": - linear._ssm_executor.raw_keys.zero_() - linear._ssm_executor.raw_values.zero_() - linear._ssm_executor.gates.fill_(-0.5) - linear._ssm_executor.betas.zero_() - linear._ssm_executor.cursors[0] = 1 + linear.raw_keys.zero_() + linear.raw_values.zero_() + linear.gates.fill_(-0.5) + linear.betas.zero_() + linear.cursors[0] = 1 ssm = linear.accepted_ssm(0).clone() - linear._ssm_executor.cursors[2] = 3 + linear.cursors[2] = 3 if mtp_step: linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( @@ -175,16 +175,20 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - assert torch.equal(linear.ssm_state[:, 2 * linear.ssm_slots_per_req], ssm) + restored_slot = 2 if mode == "replay" else 2 * linear.ssm_slots_per_req + assert torch.equal(linear.ssm_state[:, restored_slot], ssm) if mode == "replay": - assert linear._ssm_executor.cursors[2] == 0 + assert linear.cursors[2] == 0 if mtp_step: assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 assert torch.equal(tail.buffer, unchanged) resources.reset_bounded_att_state([2]) assert not linear.conv_state[:, 2].any() - assert not linear.ssm_state[:, 2 * linear.ssm_slots_per_req : 3 * linear.ssm_slots_per_req].any() + if mode == "replay": + assert not linear.ssm_state[:, 2].any() + else: + 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_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 704849ecba..bc6c1428cb 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -293,7 +293,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod if mtp_step: accepted = torch.zeros(4, dtype=torch.int32, device="cuda") accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 - selected._ssm_executor.accept_updates(ids, accepted) + selected.accept_updates(ids, accepted) for req, length in zip(requests, lengths): torch.testing.assert_close( selected.accepted_ssm(req)[0], diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index c76091af1e..5656b4bf53 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -126,7 +126,7 @@ def test_pd_kv_and_runtime_state_roundtrip( ): 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 + dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.ssm_slots_per_req mtp_size = get_env_start_args().mtp_step + 1 layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 @@ -165,9 +165,10 @@ 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 * mem.linear.ssm_slots_per_req].copy_(ssm[:, heads]) + src_slot = src_req if source_mode == "replay" else src_req * mem.linear.ssm_slots_per_req + mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) if source_mode == "replay": - backend = mem.linear._ssm_executor + backend = mem.linear backend.raw_keys.zero_() backend.raw_values.zero_() backend.gates.fill_(-0.5) @@ -188,7 +189,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() if dest_mode == "replay": - mem.linear._ssm_executor.cursors[dst_req] = 3 + mem.linear.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: @@ -260,7 +261,7 @@ def test_pd_kv_and_runtime_state_roundtrip( 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_executor.cursors[dst_req] == 0 + assert mem.linear.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( From 3f9f313a12cb81016aab4493a2d91f29eb2aef99 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:02:30 +0000 Subject: [PATCH 12/54] refactor: drop the SSM slot property and keyword-only arguments Native slot stride is mtp_step + 1 at each call site. New GDN and KDA constructors take their replay arguments as ordinary parameters. Co-authored-by: Cursor --- .../bounded_att/gdn/create_gdn_state.py | 1 - .../attention_state/bounded_att/gdn/gdn_state.py | 6 ------ .../bounded_att/gdn/replay_gdn_state.py | 7 ++----- .../bounded_att/kda/create_kda_state.py | 1 - .../attention_state/bounded_att/kda/kda_state.py | 6 ------ .../bounded_att/kda/replay_kda_state.py | 7 ++----- .../att_resources/test_hybrid_layer_mapping.py | 3 ++- .../common/test_hybrid_req_manager_pd_sharing.py | 6 +++--- unit_tests/common/test_linear_state.py | 14 +++++++------- unit_tests/models/glm5_next/test_cache.py | 4 ++-- unit_tests/models/glm5_next/test_pd_cache.py | 4 ++-- 11 files changed, 20 insertions(+), 39 deletions(-) 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 dcf92f65c4..82a07fc486 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -13,7 +13,6 @@ 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, 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 95cdef3697..b8a641a73d 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -71,12 +71,6 @@ def __init__( self.big_pages_pool = None self.small_pages_pool = None - @property - def ssm_slots_per_req(self) -> int: - """每个请求的物理 SSM 槽位数，包含 MTP 候选。""" - - return self.mtp_step + 1 - 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)}" diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index ab4b2c9a55..2f4183fa1f 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -135,7 +135,6 @@ def __init__( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", - *, ssm_state_mode: str = "replay", ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, @@ -196,7 +195,6 @@ def from_ssm( capacity: int, verify_width: int, activation_dtype: torch.dtype = torch.bfloat16, - *, num_key_heads: int | None = None, kda: bool = False, gate_lower_bound: float = -5.0, @@ -362,7 +360,7 @@ def read_from_pd_page( ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, *, fold: bool) -> None: + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" if self.config is None or not self.config.logical_layer_ids: @@ -514,7 +512,6 @@ def export_accepted_state( req_indices: list[int], targets: list[torch.Tensor], accepted_indices: torch.Tensor | None = None, - *, flush: bool = False, ) -> None: """将多个 accepted 快照直接写入调用方目标页。""" @@ -547,7 +544,7 @@ def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = Non self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" if not req_indices: 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 f94b57501d..1742441389 100644 --- a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -13,7 +13,6 @@ 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, 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 5bb9fc780d..e5480a3ab0 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -71,12 +71,6 @@ def __init__( self.big_pages_pool = None self.small_pages_pool = None - @property - def ssm_slots_per_req(self) -> int: - """每个请求的物理 SSM 槽位数，包含 MTP 候选。""" - - return self.mtp_step + 1 - 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)}" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 56d573eef6..94bb4741a9 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -135,7 +135,6 @@ def __init__( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", - *, ssm_state_mode: str = "replay", ssm_replay_config: dict | None = None, gate_lower_bound: float = -5.0, @@ -196,7 +195,6 @@ def from_ssm( capacity: int, verify_width: int, activation_dtype: torch.dtype = torch.bfloat16, - *, num_key_heads: int | None = None, kda: bool = True, gate_lower_bound: float = -5.0, @@ -362,7 +360,7 @@ def read_from_pd_page( ) -> None: read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, *, fold: bool) -> None: + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" if self.config is None or not self.config.logical_layer_ids: @@ -514,7 +512,6 @@ def export_accepted_state( req_indices: list[int], targets: list[torch.Tensor], accepted_indices: torch.Tensor | None = None, - *, flush: bool = False, ) -> None: """将多个 accepted 快照直接写入调用方目标页。""" @@ -547,7 +544,7 @@ def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = Non self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], *, flush: bool = False) -> None: + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" if not req_indices: 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 2b98f3da97..fb9ce7ba5b 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -79,7 +79,8 @@ def test_hybrid_layer_mapping_and_mtp_states( device="meta", ) if mode == "native": - assert bounded_state.ssm_slots_per_req == 3 + stride = bounded_state.mtp_step + 1 + assert bounded_state.ssm_state.shape[1] == bounded_state.request_capacity * stride if mode == "replay": assert bounded_state.capacity == 8 manager = AttentionStateManager(request_capacity=1) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index 807aac36c2..47beee1eef 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -28,7 +28,7 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): actual = state.get_accepted_state(2) else: width = state.config.conv_width - actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * state.ssm_slots_per_req]) + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * (state.mtp_step + 1)]) 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 @@ -115,10 +115,10 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( width = state.config.conv_width expected = [ state.conv_state[:, 1, :, :width].cpu().tolist(), - state.ssm_state[:, state.ssm_slots_per_req].cpu().tolist(), + state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), ] before = vars(state).copy() - source_slot = 1 if mode == "replay" else state.ssm_slots_per_req + source_slot = 1 if mode == "replay" else state.mtp_step + 1 source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() snapshots = [] for pool in (state.big_pages_pool, state.small_pages_pool): diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index b929b40afd..a3a11d0b7b 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -51,7 +51,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( state.mtp_state_index[src] = offset if mode == "native": expected_conv = state.conv_state[:, src, ..., :3].clone() - expected_ssm = state.ssm_state[:, src * state.ssm_slots_per_req].clone() + expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() else: expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() expected_ssm = state.ssm_state[:, src].clone() @@ -90,7 +90,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( actual_conv, actual_ssm = actual.conv, actual.ssm else: actual_conv = state.conv_state[:, dst, ..., :3] - actual_ssm = state.ssm_state[:, dst * state.ssm_slots_per_req] + actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) if state.mtp_state_index is not None: @@ -102,8 +102,8 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if mode == "replay": assert not state.ssm_state[:, dst].any() else: - start = dst * state.ssm_slots_per_req - assert not state.ssm_state[:, start : start + state.ssm_slots_per_req].any() + start = dst * (state.mtp_step + 1) + assert not state.ssm_state[:, start : start + state.mtp_step + 1].any() def copy_on_cpu(*, copies): @@ -136,7 +136,7 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty expected = [ ( state.conv_state[:, req, :, : state.config.conv_width].clone(), - state.ssm_state[:, req * state.ssm_slots_per_req].clone(), + state.ssm_state[:, req * (state.mtp_step + 1)].clone(), ) for req in requests ] @@ -173,7 +173,7 @@ def no_snapshot(*args): actual_conv, actual_ssm = actual.conv, actual.ssm else: actual_conv = target.conv_state[:, req, :, : target.config.conv_width] - actual_ssm = target.ssm_state[:, req * target.ssm_slots_per_req] + actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) assert torch.equal(state.ssm_state, before) @@ -233,7 +233,7 @@ def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(mo state.conv_state[:, src, :2, :3].fill_(10 + rank) state.conv_state[:, src, 2:4, :3].fill_(20 + rank) state.conv_state[:, src, 4:, :3].fill_(30 + rank) - state.ssm_state[:, src * state.ssm_slots_per_req].fill_(40 + rank) + state.ssm_state[:, src * (state.mtp_step + 1)].fill_(40 + rank) states.append(state) layouts.append(layout) snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 36431fb3cd..f583bcf7da 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -175,7 +175,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - restored_slot = 2 if mode == "replay" else 2 * linear.ssm_slots_per_req + restored_slot = 2 if mode == "replay" else 2 * (linear.mtp_step + 1) assert torch.equal(linear.ssm_state[:, restored_slot], ssm) if mode == "replay": assert linear.cursors[2] == 0 @@ -188,7 +188,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( if mode == "replay": assert not linear.ssm_state[:, 2].any() else: - assert not linear.ssm_state[:, 2 * linear.ssm_slots_per_req : 3 * linear.ssm_slots_per_req].any() + assert not linear.ssm_state[:, 2 * (linear.mtp_step + 1) : 3 * (linear.mtp_step + 1)].any() assert not tail.buffer[:, 2, :, : tail.config.candidate_dim].any() assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim :]).all() assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 5656b4bf53..7dbe235831 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -126,7 +126,7 @@ def test_pd_kv_and_runtime_state_roundtrip( ): source = make_mems(prefill_tp, mode=source_mode) dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) - dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.ssm_slots_per_req + dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.mtp_step + 1 mtp_size = get_env_start_args().mtp_step + 1 layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 @@ -165,7 +165,7 @@ 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)) - src_slot = src_req if source_mode == "replay" else src_req * mem.linear.ssm_slots_per_req + src_slot = src_req if source_mode == "replay" else src_req * (mem.linear.mtp_step + 1) mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) if source_mode == "replay": backend = mem.linear From 8d6308c844bd35fca56e6149f28f66855590530c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:17:16 +0000 Subject: [PATCH 13/54] refactor: keep replay page layout on the state GDN and KDA no longer split page copies and accepted snapshots into separate modules. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_checkpoint.py | 296 ------------ .../bounded_att/gdn/replay_gdn_snapshot.py | 115 ----- .../bounded_att/gdn/replay_gdn_state.py | 420 +++++++++++++++--- .../bounded_att/kda/replay_kda_checkpoint.py | 296 ------------ .../bounded_att/kda/replay_kda_snapshot.py | 115 ----- .../bounded_att/kda/replay_kda_state.py | 419 ++++++++++++++--- 6 files changed, 717 insertions(+), 944 deletions(-) delete mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py delete mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py delete mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py delete mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py deleted file mode 100644 index 99f197b40b..0000000000 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_checkpoint.py +++ /dev/null @@ -1,296 +0,0 @@ -"""GDN replay 的页布局和快照拷贝。 - -这里只处理 Conv/SSM 快照页，不读取 replay cursor 或 history。 -""" - -from __future__ import annotations - -from math import prod -from typing import List - -import torch - -from lightllm.utils.log_utils import init_logger - -from ...types import ( - AxisMapping, - PageKind, - StateTransferLayout, - TensorField, - TransferRegion, -) -from ..copy_kernel import copy_bounded_att_tensors -from ..transfer_utils import bounded_att_buffer_copies -from .gdn_pages_pool import GDNPagesPool - -logger = init_logger(__name__) - - -def build_transfer_regions(state) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = state.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", - state.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", - state.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)", - state.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - - -def page_pool(state, page_kind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {state.state_name}" - return pool - - -def tensor_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")), - } - - -def write_pages_to_cpu( - state, - 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 == state.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 = page_pool(state, 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(tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - -def read_pages_from_cpu( - state, - 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 == state.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 = state.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 = [tensor_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( - state, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state, -) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = state.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 init_page_pools( - state, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, -) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert state.big_pages_pool is None and state.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - state.big_pages_pool = ( - GDNPagesPool( - size=big_page_capacity, - config=state.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - state.small_pages_pool = ( - GDNPagesPool( - size=small_page_capacity, - config=state.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - -# ======================================================================== -# 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 -# ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py deleted file mode 100644 index f6492f34b9..0000000000 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_snapshot.py +++ /dev/null @@ -1,115 +0,0 @@ -"""GDN replay 的已接受快照读写。 - -页布局在 replay_gdn_checkpoint。这里决定写哪一段 Conv、是否物化 SSM，以及何时清 history。 -""" - -from __future__ import annotations - -import torch - -from ..copy_kernel import copy_bounded_att_tensors -from ..transfer_utils import bounded_att_buffer_copies -from .replay_gdn_checkpoint import page_pool, tensor_fields - - -def accepted_conv(state, req_idx: int) -> torch.Tensor: - """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" - width = state.config.conv_width - conv = state.conv_state[:, req_idx] - if state.mtp_state_index is None: - return conv[:, :, :width] - offset = state.mtp_state_index[req_idx] - return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - - -def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: - """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = page_pool(state, page_kind) - copy_bounded_att_tensors( - copies=[ - (accepted_conv(state, req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - ) - state.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) - - -def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: - """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = page_pool(state, page_kind) - width = state.config.conv_width - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - state.clear_history(req_idx) - conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) - ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if state.mtp_state_index is not None: - for req_idx in req_indices: - state.mtp_state_index[req_idx].fill_(0) - - -def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_bounded_state) -> None: - """复制已接受运行态到目标单槽，并清空目标 history。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - width = state.config.conv_width - conv_copies = [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - state.clear_history(dst_req) - conv_copies.append( - ( - accepted_conv(remote_dp_bounded_state, src_req), - state.conv_state[:, dst_req, :, :width], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.export_ssm(src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices]) - if state.mtp_state_index is not None: - for dst_req in dst_req_indices: - state.mtp_state_index[dst_req].fill_(0) - - -def write_pd(state, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" - assert layout.state_name == state.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = state.get_accepted_state(req_idx) - fields = tensor_fields(accepted.conv, accepted.ssm) - 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_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): - """读入基准 Conv/SSM，并清空该请求的 history。""" - assert layout.state_name == state.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - state.clear_history(req_idx) - fields = tensor_fields( - state.conv_state[:, req_idx, :, : state.config.conv_width], - state.ssm_state[:, req_idx], - ) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if state.mtp_state_index is not None: - state.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 2f4183fa1f..4727f6a209 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -1,13 +1,14 @@ """GDN replay 的 bounded attention 状态。 -Conv、canonical SSM 和 replay history 都定义在本类。页布局在 replay_gdn_checkpoint， -已接受快照在 replay_gdn_snapshot。一层 decode 重放在 replay_gdn_step。 +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_gdn_step。 """ from __future__ import annotations from collections.abc import Sequence from functools import lru_cache +from math import prod from typing import NamedTuple import torch @@ -20,24 +21,17 @@ _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig -from .replay_gdn_checkpoint import ( - build_transfer_regions, - copy_big_pages, - init_page_pools, - read_pages_from_cpu, - write_pages_to_cpu, -) -from .replay_gdn_snapshot import ( - accepted_conv, - copy_runtime_from_dp, - read_pd, - restore_runtime, - save_runtime, - write_pd, -) +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .gdn_pages_pool import GDNPagesPool + +logger = init_logger(__name__) class _AcceptedTensors(NamedTuple): @@ -108,6 +102,10 @@ def _bind_history( class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + _run_config: dict[str, int] | None = None @property @@ -243,9 +241,36 @@ def initialize_page_pools( device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: - init_page_pools(self, big_page_capacity, small_page_capacity, device, pin_memory) + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + GDNPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + GDNPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" return logical_layer_id in self._logical_to_storage_layer def storage_layer(self, logical_layer_id: int) -> int: @@ -257,41 +282,6 @@ def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self): - return build_transfer_regions(self) - - def write_state_pages_to_cpu( - self, - page_kind, - page_slot_ids, - cpu_page_buffers, - cpu_page_indices, - layout, - cpu_page_ready_flags=None, - ): - write_pages_to_cpu( - self, - page_kind, - page_slot_ids, - cpu_page_buffers, - cpu_page_indices, - layout, - cpu_page_ready_flags, - ) - - def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): - read_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout) - - def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): - copy_big_pages(self, src_page_ids, dst_page_ids, remote_dp_bounded_state) - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self): - return build_transfer_regions(self) - def reset_bounded_att_state(self, req_indices: list[int]) -> None: """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" assert isinstance(req_indices, list), "request indices must be a list" @@ -309,7 +299,17 @@ def save_bounded_att_runtime_to_pages( page_kind, page_slot_ids: list[int], ) -> None: - save_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) def restore_bounded_att_runtime_from_pages( self, @@ -318,7 +318,21 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + 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: """提交已接受前缀，折进 canonical SSM，并清空 history。""" @@ -330,7 +344,102 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(accepted_conv(self, req_idx), self.accepted_ssm(req_idx)) + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def copy_runtime_state_from_dp( self, @@ -338,7 +447,37 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayGDNState", ) -> None: - copy_runtime_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state) + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + 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._transfer_regions() def write_to_pd_page( self, @@ -346,9 +485,21 @@ def write_to_pd_page( pd_page_buffers: torch.Tensor, pd_page_index: int, state_token_len: int, - layout, + layout: StateTransferLayout, ) -> None: - write_pd(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout) + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + 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, @@ -356,9 +507,156 @@ def read_from_pd_page( req_idx: int, pd_page_index: int, state_token_len: int, - layout, + layout: StateTransferLayout, ) -> None: - read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、页布局 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + # ======================================================================== + # 七、Replay history + # ======================================================================== def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py deleted file mode 100644 index 5723b71db9..0000000000 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_checkpoint.py +++ /dev/null @@ -1,296 +0,0 @@ -"""KDA replay 的页布局和快照拷贝。 - -这里只处理 Conv/SSM 快照页，不读取 replay cursor 或 history。 -""" - -from __future__ import annotations - -from math import prod -from typing import List - -import torch - -from lightllm.utils.log_utils import init_logger - -from ...types import ( - AxisMapping, - PageKind, - StateTransferLayout, - TensorField, - TransferRegion, -) -from ..copy_kernel import copy_bounded_att_tensors -from ..transfer_utils import bounded_att_buffer_copies -from .kda_pages_pool import KDAPagesPool - -logger = init_logger(__name__) - - -def build_transfer_regions(state) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = state.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", - state.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", - state.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)", - state.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - - -def page_pool(state, page_kind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {state.state_name}" - return pool - - -def tensor_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")), - } - - -def write_pages_to_cpu( - state, - 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 == state.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 = page_pool(state, 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(tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - -def read_pages_from_cpu( - state, - 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 == state.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 = state.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 = [tensor_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( - state, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state, -) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = state.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 init_page_pools( - state, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, -) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert state.big_pages_pool is None and state.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - state.big_pages_pool = ( - KDAPagesPool( - size=big_page_capacity, - config=state.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - state.small_pages_pool = ( - KDAPagesPool( - size=small_page_capacity, - config=state.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - -# ======================================================================== -# 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 -# ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py deleted file mode 100644 index c4d117826b..0000000000 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_snapshot.py +++ /dev/null @@ -1,115 +0,0 @@ -"""KDA replay 的已接受快照读写。 - -页布局在 replay_kda_checkpoint。这里决定写哪一段 Conv、是否物化 SSM，以及何时清 history。 -""" - -from __future__ import annotations - -import torch - -from ..copy_kernel import copy_bounded_att_tensors -from ..transfer_utils import bounded_att_buffer_copies -from .replay_kda_checkpoint import page_pool, tensor_fields - - -def accepted_conv(state, req_idx: int) -> torch.Tensor: - """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" - width = state.config.conv_width - conv = state.conv_state[:, req_idx] - if state.mtp_state_index is None: - return conv[:, :, :width] - offset = state.mtp_state_index[req_idx] - return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - - -def save_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: - """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = page_pool(state, page_kind) - copy_bounded_att_tensors( - copies=[ - (accepted_conv(state, req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - ) - state.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) - - -def restore_runtime(state, req_indices, state_token_lens, page_kind, page_slot_ids) -> None: - """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = page_pool(state, page_kind) - width = state.config.conv_width - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - state.clear_history(req_idx) - conv_copies.append((pool.conv_state[slot_id], state.conv_state[:, req_idx, :, :width])) - ssm_copies.append((pool.ssm_state[slot_id], state.ssm_state[:, req_idx])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if state.mtp_state_index is not None: - for req_idx in req_indices: - state.mtp_state_index[req_idx].fill_(0) - - -def copy_runtime_from_dp(state, src_req_indices, dst_req_indices, remote_dp_bounded_state) -> None: - """复制已接受运行态到目标单槽，并清空目标 history。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(state) and remote_dp_bounded_state.config == state.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - width = state.config.conv_width - conv_copies = [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - state.clear_history(dst_req) - conv_copies.append( - ( - accepted_conv(remote_dp_bounded_state, src_req), - state.conv_state[:, dst_req, :, :width], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.export_ssm(src_req_indices, [state.ssm_state[:, dst_req] for dst_req in dst_req_indices]) - if state.mtp_state_index is not None: - for dst_req in dst_req_indices: - state.mtp_state_index[dst_req].fill_(0) - - -def write_pd(state, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" - assert layout.state_name == state.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = state.get_accepted_state(req_idx) - fields = tensor_fields(accepted.conv, accepted.ssm) - 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_pd(state, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): - """读入基准 Conv/SSM，并清空该请求的 history。""" - assert layout.state_name == state.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - state.clear_history(req_idx) - fields = tensor_fields( - state.conv_state[:, req_idx, :, : state.config.conv_width], - state.ssm_state[:, req_idx], - ) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if state.mtp_state_index is not None: - state.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 94bb4741a9..0ac40f46de 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -1,13 +1,14 @@ """KDA replay 的 bounded attention 状态。 -Conv、canonical SSM 和 replay history 都定义在本类。页布局在 replay_kda_checkpoint， -已接受快照在 replay_kda_snapshot。一层 decode 重放在 replay_kda_step。 +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_kda_step。 """ from __future__ import annotations from collections.abc import Sequence from functools import lru_cache +from math import prod from typing import NamedTuple import torch @@ -20,24 +21,17 @@ _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState from ..conv_ssm_types import ConvSSMStateConfig -from .replay_kda_checkpoint import ( - build_transfer_regions, - copy_big_pages, - init_page_pools, - read_pages_from_cpu, - write_pages_to_cpu, -) -from .replay_kda_snapshot import ( - accepted_conv, - copy_runtime_from_dp, - read_pd, - restore_runtime, - save_runtime, - write_pd, -) +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .kda_pages_pool import KDAPagesPool + +logger = init_logger(__name__) class _AcceptedTensors(NamedTuple): @@ -108,6 +102,10 @@ def _bind_history( class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + _run_config: dict[str, int] | None = None @property @@ -243,9 +241,36 @@ def initialize_page_pools( device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: - init_page_pools(self, big_page_capacity, small_page_capacity, device, pin_memory) + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + KDAPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + KDAPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" return logical_layer_id in self._logical_to_storage_layer def storage_layer(self, logical_layer_id: int) -> int: @@ -257,41 +282,6 @@ def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self): - return build_transfer_regions(self) - - def write_state_pages_to_cpu( - self, - page_kind, - page_slot_ids, - cpu_page_buffers, - cpu_page_indices, - layout, - cpu_page_ready_flags=None, - ): - write_pages_to_cpu( - self, - page_kind, - page_slot_ids, - cpu_page_buffers, - cpu_page_indices, - layout, - cpu_page_ready_flags, - ) - - def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): - read_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout) - - def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): - copy_big_pages(self, src_page_ids, dst_page_ids, remote_dp_bounded_state) - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self): - return build_transfer_regions(self) - def reset_bounded_att_state(self, req_indices: list[int]) -> None: """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" assert isinstance(req_indices, list), "request indices must be a list" @@ -309,7 +299,17 @@ def save_bounded_att_runtime_to_pages( page_kind, page_slot_ids: list[int], ) -> None: - save_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) def restore_bounded_att_runtime_from_pages( self, @@ -318,7 +318,21 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - restore_runtime(self, req_indices, state_token_lens, page_kind, page_slot_ids) + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + 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: """提交已接受前缀，折进 canonical SSM，并清空 history。""" @@ -330,7 +344,102 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(accepted_conv(self, req_idx), self.accepted_ssm(req_idx)) + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def copy_runtime_state_from_dp( self, @@ -338,7 +447,37 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayKDAState", ) -> None: - copy_runtime_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state) + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + 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._transfer_regions() def write_to_pd_page( self, @@ -346,9 +485,21 @@ def write_to_pd_page( pd_page_buffers: torch.Tensor, pd_page_index: int, state_token_len: int, - layout, + layout: StateTransferLayout, ) -> None: - write_pd(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout) + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + 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, @@ -356,10 +507,156 @@ def read_from_pd_page( req_idx: int, pd_page_index: int, state_token_len: int, - layout, + layout: StateTransferLayout, ) -> None: - read_pd(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout) + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、页布局 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + # ======================================================================== + # 七、Replay history + # ======================================================================== def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" From c3075fe966759d6d54da71f223a2fb62cabe7429 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:26:31 +0000 Subject: [PATCH 14/54] refactor: let replay state commit accepted lengths directly The shared acceptance wrapper only unpacked tensors the state already receives. Private helpers stay at the end of the replay state. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 447 +++++++++--------- .../bounded_att/kda/replay_kda_state.py | 445 +++++++++-------- lightllm/common/attention_state/commit.py | 70 --- .../linear_att/test_acceptance_fusion.py | 7 +- .../linear_att/test_replayssm.py | 4 +- .../linear_att/test_replayssm_execution.py | 7 +- 6 files changed, 448 insertions(+), 532 deletions(-) delete mode 100644 lightllm/common/attention_state/commit.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 4727f6a209..353fb686ec 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -39,66 +39,6 @@ class _AcceptedTensors(NamedTuple): ssm: torch.Tensor -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode - - -def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -def _bind_history( - owner, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype, - num_key_heads: int | None, - gate_lower_bound: float, - compute_mode: str, - run_config: dict[str, int] | None, -) -> None: - """在 owner 上创建 GDN replay 的双半区 history 和 cursor。""" - - assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - owner.ssm_state = ssm - owner.kda = False - owner.gate_lower_bound = gate_lower_bound - owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 - owner.capacity = capacity - owner.verify_width = verify_width - owner.compute_mode = compute_mode - owner.run_config = run_config - owner._config_is_fixed = run_config is not None - owner.batch_configs = {} - layers, slots, hv, k, v = ssm.shape - owner.num_key_heads = hv if num_key_heads is None else num_key_heads - assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 - owner.hold_req_id = slots - 1 - owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.forward_cursors = torch.zeros_like(owner.cursors) - # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) - owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) - owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) - - class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -529,7 +469,175 @@ def read_from_pd_page( self.mtp_state_index[req_idx].fill_(0) # ======================================================================== - # 六、页布局 + # 六、Replay history + # ======================================================================== + + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: + """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + + if self.config is None or not self.config.logical_layer_ids: + return + if logical_layer_id != self.config.logical_layer_ids[-1]: + return + if fold: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.ssm_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, + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.ssm_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.ssm_state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + False, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor，保留每个请求本轮已接受的前缀。""" + + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """非 MTP decode 把单 token 记入当前 history。""" + + self.accept_updates(reqs) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把已接受 history 折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + self._materialize_accepted( + _upload_indices(req_indices, self.ssm_state.device, torch.int32), + _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), + dtype=self.ssm_state.dtype, + device=self.ssm_state.device, + ) + self._materialize_accepted( + torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + # ======================================================================== + # 七、私有辅助实现 # ======================================================================== def _transfer_regions(self) -> tuple[TransferRegion, ...]: @@ -654,108 +762,6 @@ def _accepted_conv(self, req_idx: int) -> torch.Tensor: offset = self.mtp_state_index[req_idx] return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - # ======================================================================== - # 七、Replay history - # ======================================================================== - - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" - - if self.config is None or not self.config.logical_layer_ids: - return - if logical_layer_id != self.config.logical_layer_ids[-1]: - return - if fold: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) - - def clear_history(self, req: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM。""" - - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - - positions = torch.empty_like(reqs) - layers, _, hv, kd, vd = self.ssm_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, - num_warps=4, - ) - programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) - _fold_history[(programs,)]( - self.ssm_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.ssm_state.stride(0), - self.raw_keys.stride(0), - self.raw_values.stride(0), - self.gates.stride(0), - self.betas.stride(0), - False, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor，保留每个请求本轮已接受的前缀。""" - - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - reqs, - accepted_state_indices, - reqs.numel(), - self.hold_req_id, - self.verify_width, - 256, - ) - - def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 把单 token 记入当前 history。""" - - self.accept_updates(reqs) - - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把已接受 history 折回 canonical SSM，并清空 cursor。""" - - if reqs.numel() == 0: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False) - self.cursors.index_fill_(0, reqs.long(), 0) - def _materialize_accepted( self, reqs: torch.Tensor, @@ -796,73 +802,62 @@ def _materialize_accepted( num_stages=config["num_stages"], ) - def accepted_ssm(self, req_idx: int) -> torch.Tensor: - """返回不修改 canonical SSM 和 history 的单请求快照。""" - - layers, _, hv, k, v = self.ssm_state.shape - output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) - reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) - self._materialize_accepted(reqs, output, snapshot=True) - return output - - def export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页。""" - if not req_indices: - return - shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) - self._materialize_accepted( - _upload_indices(req_indices, self.ssm_state.device, torch.int32), - _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) +def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """从 acceptance 取出请求和接受长度，供直接持有 history 的测试调用。""" + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, +) -> None: + """在 owner 上创建 GDN replay 的双半区 history 和 cursor。""" - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.ssm_state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), - dtype=self.ssm_state.dtype, - device=self.ssm_state.device, - ) - self._materialize_accepted( - torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.ssm_state[:, req].copy_(snap) - self.clear_history(req) + assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = False + owner.gate_lower_bound = gate_lower_bound + owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.forward_cursors = torch.zeros_like(owner.cursors) + # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 0ac40f46de..06d21e14a0 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -39,66 +39,6 @@ class _AcceptedTensors(NamedTuple): ssm: torch.Tensor -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode - - -def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - -def _bind_history( - owner, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype, - num_key_heads: int | None, - gate_lower_bound: float, - compute_mode: str, - run_config: dict[str, int] | None, -) -> None: - """在 owner 上创建 KDA replay 的双半区 history 和 cursor。""" - - assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - owner.ssm_state = ssm - owner.kda = True - owner.gate_lower_bound = gate_lower_bound - owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 - owner.capacity = capacity - owner.verify_width = verify_width - owner.compute_mode = compute_mode - owner.run_config = run_config - owner._config_is_fixed = run_config is not None - owner.batch_configs = {} - layers, slots, hv, k, v = ssm.shape - owner.num_key_heads = hv if num_key_heads is None else num_key_heads - assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 - owner.hold_req_id = slots - 1 - owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.forward_cursors = torch.zeros_like(owner.cursors) - # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) - owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=torch.float32, device=ssm.device) - owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) - - class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -529,7 +469,174 @@ def read_from_pd_page( self.mtp_state_index[req_idx].fill_(0) # ======================================================================== - # 六、页布局 + # 六、Replay history + # ======================================================================== + def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: + """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + + if self.config is None or not self.config.logical_layer_ids: + return + if logical_layer_id != self.config.logical_layer_ids[-1]: + return + if fold: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.ssm_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, + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.ssm_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.ssm_state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + True, + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: + """更新 cursor，保留每个请求本轮已接受的前缀。""" + + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """非 MTP decode 把单 token 记入当前 history。""" + + self.accept_updates(reqs) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把已接受 history 折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + accepted_indices: torch.Tensor | None = None, + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + self._materialize_accepted( + _upload_indices(req_indices, self.ssm_state.device, torch.int32), + _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = torch.empty( + (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), + dtype=self.ssm_state.dtype, + device=self.ssm_state.device, + ) + self._materialize_accepted( + torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), + snap, + snapshot=True, + ) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + # ======================================================================== + # 七、私有辅助实现 # ======================================================================== def _transfer_regions(self) -> tuple[TransferRegion, ...]: @@ -654,107 +761,6 @@ def _accepted_conv(self, req_idx: int) -> torch.Tensor: offset = self.mtp_state_index[req_idx] return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - # ======================================================================== - # 七、Replay history - # ======================================================================== - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" - - if self.config is None or not self.config.logical_layer_ids: - return - if logical_layer_id != self.config.logical_layer_ids[-1]: - return - if fold: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) - - def clear_history(self, req: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM。""" - - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" - - positions = torch.empty_like(reqs) - layers, _, hv, kd, vd = self.ssm_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, - num_warps=4, - ) - programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) - _fold_history[(programs,)]( - self.ssm_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.ssm_state.stride(0), - self.raw_keys.stride(0), - self.raw_values.stride(0), - self.gates.stride(0), - self.betas.stride(0), - True, - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor，保留每个请求本轮已接受的前缀。""" - - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - reqs, - accepted_state_indices, - reqs.numel(), - self.hold_req_id, - self.verify_width, - 256, - ) - - def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 把单 token 记入当前 history。""" - - self.accept_updates(reqs) - - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把已接受 history 折回 canonical SSM，并清空 cursor。""" - - if reqs.numel() == 0: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False) - self.cursors.index_fill_(0, reqs.long(), 0) - def _materialize_accepted( self, reqs: torch.Tensor, @@ -795,73 +801,62 @@ def _materialize_accepted( num_stages=config["num_stages"], ) - def accepted_ssm(self, req_idx: int) -> torch.Tensor: - """返回不修改 canonical SSM 和 history 的单请求快照。""" - layers, _, hv, k, v = self.ssm_state.shape - output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) - reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) - self._materialize_accepted(reqs, output, snapshot=True) - return output - - def export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页。""" - - if not req_indices: - return - shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) - self._materialize_accepted( - _upload_indices(req_indices, self.ssm_state.device, torch.int32), - _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: + config = dict(config or {}) + capacity = int(config.get("cache_len", 16)) + compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + return capacity, compute_mode - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) +def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: + """把 Python 索引上传到目标 device。""" - def commit(self, acceptance, indices: torch.Tensor, state_name: str | None = None) -> None: - """从 acceptance 取出请求和接受长度，供直接持有 history 的测试调用。""" + return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - self.commit_accepted(indices, acceptance.req_indices, acceptance.accepted_lengths) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, +) -> None: + """在 owner 上创建 KDA replay 的双半区 history 和 cursor。""" - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.ssm_state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), - dtype=self.ssm_state.dtype, - device=self.ssm_state.device, - ) - self._materialize_accepted( - torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), - snap, - snapshot=True, - ) - dst.copy_(snap) - if flush: - self.ssm_state[:, req].copy_(snap) - self.clear_history(req) + assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = True + owner.gate_lower_bound = gate_lower_bound + owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.forward_cursors = torch.zeros_like(owner.cursors) + # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py deleted file mode 100644 index f9dc147975..0000000000 --- a/lightllm/common/attention_state/commit.py +++ /dev/null @@ -1,70 +0,0 @@ -from collections.abc import Iterator -from dataclasses import dataclass - -import torch - - -@dataclass(frozen=True) -class StateCommit: - """保存一次 forward 产生的状态提交配置。""" - - state_name: str - payload: object - - -@dataclass(frozen=True) -class StateCommitBatch: - """描述一个 microbatch 分区及其各状态的 forward 配置。""" - - request_count: int - commits: tuple[StateCommit, ...] = () - - def __post_init__(self) -> None: - 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 StateAcceptance: - """描述一次 speculative verify 的接受结果。 - - ``req_indices`` 和 ``accepted_lengths`` 按未 padding 的 forward 顺序配对， - 接受长度包含基准输入 token。``batches`` 只保存相同顺序下的 microbatch - 配置，提交时由状态执行器按分区恢复产生本次结果所用的 kernel layout。 - """ - - # 接受结果是状态提交的唯一输入；batches 只携带对应 forward 的执行配置，避免调用方 - # 重新拼装请求顺序、接受长度和 microbatch 分区。 - - req_indices: torch.Tensor - accepted_lengths: torch.Tensor - batches: tuple[StateCommitBatch, ...] = () - - def validate(self, state_names: set[str] | None = None) -> None: - """检查请求、接受长度和 forward 分区是否能组成一次提交。""" - - if self.req_indices.ndim != 1 or self.accepted_lengths.shape != self.req_indices.shape: - raise ValueError("request indices and accepted lengths must share one-dimensional layout") - for value in (self.req_indices, self.accepted_lengths): - if value.dtype not in (torch.int32, torch.int64) or value.device != self.req_indices.device: - raise ValueError("acceptance tensors must be integer tensors on the same device") - if self.batches and sum(batch.request_count for batch in self.batches) != self.req_indices.numel(): - 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) -> Iterator[tuple[torch.Tensor, object | None]]: - """按 forward 分区返回请求索引和指定状态的提交配置。""" - - 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 self.req_indices[start:end], payload - start = end 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 index 20b1fa6a12..fb5e9221e0 100644 --- 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 @@ -21,7 +21,6 @@ def ssm_snapshot(backend, req): from types import SimpleNamespace -from lightllm.common.attention_state.commit import StateAcceptance from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer @@ -92,7 +91,7 @@ def accept_fused(m, args): ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) lengths = cumulative[ends.long()] - cumulative[args[0].long()] - m._ssm_executor.commit(StateAcceptance(reqs, lengths), m.mtp_state_index, "main") + m._ssm_executor.commit_accepted(m.mtp_state_index, reqs, lengths) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -136,8 +135,8 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): checks += 1 class _IndexOnly: - def commit(self, acceptance, indices, state_name=None): - commit_ssm_prefix(indices, acceptance.req_indices, acceptance.accepted_lengths) + def commit_accepted(self, indices, req_indices, accepted_lengths): + commit_ssm_prefix(indices, req_indices, accepted_lengths) m._ssm_executor = _IndexOnly() accept_fused(m, args) 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 index 362b7024a2..94e583b435 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -610,8 +610,6 @@ def test_kda_tuning_candidates_match_reference(config, compute_mode): ], ) def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): - from lightllm.common.attention_state.commit import StateAcceptance - torch.manual_seed(72) initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) @@ -666,7 +664,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa ) torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) - backend.commit(StateAcceptance(reqs[:live], lengths), indices, "main") + backend.commit_accepted(indices, reqs[:live], lengths) backend.merge_accepted_updates(reqs[:live]) torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) assert torch.equal(backend.ssm_state[0, live], initial[0, live]) 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 index 3c2a81ee92..7ed9dc7e90 100644 --- 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 @@ -19,7 +19,6 @@ def ReplaySSMExecutor(*args, **kwargs): build_dynamic_mtp_linear_att_state_params, ) from types import SimpleNamespace -from lightllm.common.attention_state.commit import StateAcceptance pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -162,10 +161,10 @@ def forward_accept(index, cu, reqs): if index == 0: ssm_backend.accept_updates(reqs, accepted) else: - manager._ssm_executor.commit( - StateAcceptance(rows[starts.long()], accepted[rows[starts.long()]] + 1), + manager._ssm_executor.commit_accepted( manager.mtp_state_index, - "main", + rows[starts.long()], + accepted[rows[starts.long()]] + 1, ) return out From 4f5f3f519d42999b4d173362f3b7268751b4923e Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:31:34 +0000 Subject: [PATCH 15/54] refactor: name the replay decode kernel helpers The step modules only launch one decode replay. Decode makes that role visible next to the attention backends. Co-authored-by: Cursor --- .../attention_state/bounded_att/gdn/replay_gdn_state.py | 2 +- .../attention_state/bounded_att/kda/replay_kda_state.py | 2 +- lightllm/common/basemodel/attention/linear/replay_gdn.py | 2 +- .../linear/{replay_gdn_step.py => replay_gdn_decode.py} | 4 ++-- lightllm/common/basemodel/attention/linear/replay_kda.py | 2 +- .../linear/{replay_kda_step.py => replay_kda_decode.py} | 4 ++-- .../common/basemodel/triton_kernel/linear_att/ssm_autotune.py | 4 ++-- unit_tests/common/test_linear_state_contract.py | 2 +- 8 files changed, 11 insertions(+), 11 deletions(-) rename lightllm/common/basemodel/attention/linear/{replay_gdn_step.py => replay_gdn_decode.py} (97%) rename lightllm/common/basemodel/attention/linear/{replay_kda_step.py => replay_kda_decode.py} (97%) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 353fb686ec..778b55f22b 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -1,7 +1,7 @@ """GDN replay 的 bounded attention 状态。 Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 -一层 decode 重放在 replay_gdn_step。 +一层 decode 重放在 replay_gdn_decode。 """ from __future__ import annotations diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 06d21e14a0..26bc95b27e 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -1,7 +1,7 @@ """KDA replay 的 bounded attention 状态。 Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 -一层 decode 重放在 replay_kda_step。 +一层 decode 重放在 replay_kda_decode。 """ from __future__ import annotations diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 6cd5b38c17..5c79aff9ce 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -38,7 +38,7 @@ ) from lightllm.utils.envs_utils import get_env_start_args -from .replay_gdn_step import begin_decode, replay +from .replay_gdn_decode import begin_decode, replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py similarity index 97% rename from lightllm/common/basemodel/attention/linear/replay_gdn_step.py rename to lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 33f09c831d..5844850c8e 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -1,6 +1,6 @@ -"""GDN replay 的一层 decode。 +"""GDN replay decode 的 kernel 辅助。 -History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置并发起重放。 +History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置，并发起一层重放。 """ from __future__ import annotations diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index d3539311ba..892c7901e3 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -34,7 +34,7 @@ build_dynamic_mtp_linear_att_state_params, ) -from .replay_kda_step import begin_decode, replay +from .replay_kda_decode import begin_decode, replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_step.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py similarity index 97% rename from lightllm/common/basemodel/attention/linear/replay_kda_step.py rename to lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 770f74ee3a..cfa59e10ab 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_step.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -1,6 +1,6 @@ -"""KDA replay 的一层 decode。 +"""KDA replay decode 的 kernel 辅助。 -History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置并发起重放。 +History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置，并发起一层重放。 """ from __future__ import annotations diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 84b96964ff..a3484c0d8f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -137,9 +137,9 @@ def replay_layer(history, *args, **kwargs): from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState if isinstance(history, ReplayKDAState): - from lightllm.common.basemodel.attention.linear.replay_kda_step import replay + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay elif isinstance(history, ReplayGDNState): - from lightllm.common.basemodel.attention.linear.replay_gdn_step import replay + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay else: raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(history).__name__}") return replay(history, *args, **kwargs) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index ba2ec508db..9b4a8b031b 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -4,7 +4,7 @@ from lightllm.common.attention_state.bounded_att import KDAState from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.linear.replay_gdn_step import run_config as gdn_run_config +from lightllm.common.basemodel.attention.linear.replay_gdn_decode import run_config as gdn_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig From 91ec748d5637b79aa435fa4d9445db497f805b05 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:40:17 +0000 Subject: [PATCH 16/54] refactor: place private helpers after the public interfaces run_config is Triton launch metadata, so it no longer sits ahead of state initialization. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 52 ++-- .../bounded_att/kda/replay_kda_state.py | 52 ++-- lightllm/common/attention_state/manager.py | 12 +- .../attention/create_linear_utils.py | 20 +- .../common/basemodel/attention/linear/gdn.py | 111 +++++---- .../common/basemodel/attention/linear/kda.py | 96 ++++---- .../basemodel/attention/linear/replay_gdn.py | 224 +++++++++--------- .../basemodel/attention/linear/replay_kda.py | 104 ++++---- .../linear_att/mtp_state_params.py | 88 +++---- .../linear_att/ssm_acceptance.py | 44 ++-- .../triton_kernel/linear_att/ssm_autotune.py | 20 +- lightllm/models/glm5_next/model.py | 56 ++--- lightllm/models/qwen3next/model.py | 28 +-- 13 files changed, 452 insertions(+), 455 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 778b55f22b..0390a85e49 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -34,11 +34,6 @@ logger = init_logger(__name__) -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -46,27 +41,6 @@ class ReplayGDNState(BoundedAttentionState): # 一、资源定义与初始化 # ======================================================================== - _run_config: dict[str, int] | None = None - - @property - def run_config(self) -> dict[str, int] | None: - """返回物化和 autotune 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: dict | None) -> None: - """保存一份与调用方脱开的 kernel 配置。""" - - if config is None: - self._run_config = None - return - self._run_config = { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } - def __init__( self, config: ConvSSMStateConfig, @@ -636,6 +610,27 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.ssm_state[:, req].copy_(snap) self.clear_history(req) + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + # ======================================================================== # 七、私有辅助实现 # ======================================================================== @@ -861,3 +856,8 @@ def _bind_history( owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 26bc95b27e..5ac70f6818 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -34,11 +34,6 @@ logger = init_logger(__name__) -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -46,27 +41,6 @@ class ReplayKDAState(BoundedAttentionState): # 一、资源定义与初始化 # ======================================================================== - _run_config: dict[str, int] | None = None - - @property - def run_config(self) -> dict[str, int] | None: - """返回物化和 autotune 使用的 kernel 配置。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: dict | None) -> None: - """保存一份与调用方脱开的 kernel 配置。""" - - if config is None: - self._run_config = None - return - self._run_config = { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } - def __init__( self, config: ConvSSMStateConfig, @@ -635,6 +609,27 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.ssm_state[:, req].copy_(snap) self.clear_history(req) + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + # ======================================================================== # 七、私有辅助实现 # ======================================================================== @@ -860,3 +855,8 @@ def _bind_history( owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..6f235d3a6d 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -59,12 +59,6 @@ def __init__(self, request_capacity: int): self._dp_shared_operator = DPSharedOperator(self) self._pd_operator = PDOperator(self) - def _reset_operators(self) -> None: - """让浅拷贝后的 Operator 指向新的资源管理器副本。""" - self._cpu_cache_operator = CpuCacheOperator(self) - self._dp_shared_operator = DPSharedOperator(self) - self._pd_operator = PDOperator(self) - def get_cpu_cache_operator(self) -> CpuCacheOperator: """返回当前资源管理器唯一的 CPU cache 操作对象。""" return self._cpu_cache_operator @@ -412,6 +406,12 @@ def commit_bounded_att_state( # 七、私有辅助实现 # ======================================================================== + def _reset_operators(self) -> None: + """让浅拷贝后的 Operator 指向新的资源管理器副本。""" + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + # 私有实现：缓存创建完成后的成员集合，模型逐层计算直接访问具体 State 成员。 @lru_cache(maxsize=None) def _get_states( diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index ffd5e7eeb9..fce8e8d78f 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -15,16 +15,6 @@ } -def _linear_backend_class(native_path: str, backend_name: str): - """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" - - if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": - return load_backend_class(native_path) - from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class - - return replay_gdn_backend_class(backend_name) - - def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): args = get_env_start_args() backend_str = _get_backend_str(args.llm_prefill_att_backend, index) @@ -63,6 +53,16 @@ def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton", return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") +def _linear_backend_class(native_path: str, backend_name: str): + """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" + + if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": + return load_backend_class(native_path) + from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class + + return replay_gdn_backend_class(backend_name) + + def _get_backend_str(backend_args, index): assert index == 1, "linear attention backend index must be 1" return backend_args[index] if len(backend_args) > index else "auto" diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..95b7fc157b 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -36,8 +36,13 @@ def __init__(self, model: "TpPartBaseModel"): def get_prefill_kernel(self): pass - def _init_linear_layer_metadata(self, network_config, tp_world_size): + 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_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": + return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def _init_linear_layer_metadata(self, network_config, tp_world_size): self.mtp_step = get_env_start_args().mtp_step # Linear attention specific dimensions @@ -100,16 +105,9 @@ 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_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(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 @@ -207,7 +205,6 @@ def _gdn_prefill_kernel( @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 @@ -233,54 +230,6 @@ def init_state(self): 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 decode_att( self, q: torch.Tensor, @@ -332,6 +281,54 @@ def decode_att( ) return core_attn_out, z + 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 _gdn_decode_kernel( self, mixed_qkv: torch.Tensor, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index cd2f8afb39..6d5ddc0cb3 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -157,54 +157,6 @@ def init_state(self): 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] - def decode_att( self, q: torch.Tensor, @@ -262,3 +214,51 @@ def decode_att( num_accepted_tokens=self.b_num_accepted_tokens, ) return output + + 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] diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 5c79aff9ce..8494b39f1b 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -57,96 +57,6 @@ def replay_gdn_backend_class(name: str): raise ValueError(f"replay GDN has no backend named {name!r}") from None -def _triton_prefill_kernel(): - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - chunk_gated_delta_rule, - ) - - return chunk_gated_delta_rule - - -def _flashqla_prefill_kernel(): - from flash_qla import chunk_gated_delta_rule - - return chunk_gated_delta_rule - - -def _configure_gdn(backend, model, prefill_kernel) -> None: - """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" - - config = model.config - tp_world_size = model.tp_world_size_ - backend.mtp_step = get_env_start_args().mtp_step - backend.num_v_heads = config["linear_num_value_heads"] - backend.num_k_heads = config["linear_num_key_heads"] - backend.head_k_dim = config["linear_key_head_dim"] - backend.head_v_dim = config["linear_value_head_dim"] - backend.key_dim = backend.head_k_dim * backend.num_k_heads - backend.value_dim = backend.head_v_dim * backend.num_v_heads - backend.conv_kernel_dim = config["linear_conv_kernel_dim"] - backend.activation = config["hidden_act"] - backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size - backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size - backend.tp_num_k_heads = backend.num_k_heads // tp_world_size - backend.tp_num_v_heads = backend.num_v_heads // tp_world_size - backend.tp_key_dim = backend.key_dim // tp_world_size - backend.tp_value_dim = backend.value_dim // tp_world_size - assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads - ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} - backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) - backend.prefill_kernel = prefill_kernel - - -def _split_qkvzba(backend, mixed_qkvzba): - qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim - z_end = qkv_dim + backend.tp_value_dim - b_end = z_end + backend.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - -def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): - query, key, value = torch.split( - mixed_qkv, - [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], - dim=-1, - ) - if decode: - batch_size = mixed_qkv.shape[0] - leading = (batch_size, 1) - else: - leading = (1, query.shape[0]) - query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) - return query, key, value - - -def _replay_state(att_state, layer_num: int) -> ReplayGDNState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=ReplayGDNState, - ) - - -def _single_replay_state(att_state) -> ReplayGDNState: - role = att_state.infer_state.att_state_role - matches = [ - state - for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayGDNState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayGDNState for {role}") - return matches[0] - - class ReplayGDNTritonAttBackend(BaseAttBackend): """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" @@ -251,6 +161,28 @@ def init_state(self): self.b_ssm_buffer_idx = self.b_conv_buffer_idx begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend = self.backend + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) + return core, z + return self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state if draft_step == 0: @@ -282,28 +214,6 @@ def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None - params = att_control.linear_att_decode_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend = self.backend - mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - state = _replay_state(self, layer_num) - conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) - draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) - if draft_step > 0: - core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) - return core, z - return self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) - def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): backend = self.backend query, key, value, z, a, b = conv_pack_gdn_decode_inputs( @@ -366,3 +276,93 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step self.b1_mtp_cu_q_seq_len, ) return output + + +def _triton_prefill_kernel(): + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + chunk_gated_delta_rule, + ) + + return chunk_gated_delta_rule + + +def _flashqla_prefill_kernel(): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + +def _configure_gdn(backend, model, prefill_kernel) -> None: + """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" + + config = model.config + tp_world_size = model.tp_world_size_ + backend.mtp_step = get_env_start_args().mtp_step + backend.num_v_heads = config["linear_num_value_heads"] + backend.num_k_heads = config["linear_num_key_heads"] + backend.head_k_dim = config["linear_key_head_dim"] + backend.head_v_dim = config["linear_value_head_dim"] + backend.key_dim = backend.head_k_dim * backend.num_k_heads + backend.value_dim = backend.head_v_dim * backend.num_v_heads + backend.conv_kernel_dim = config["linear_conv_kernel_dim"] + backend.activation = config["hidden_act"] + backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size + backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size + backend.tp_num_k_heads = backend.num_k_heads // tp_world_size + backend.tp_num_v_heads = backend.num_v_heads // tp_world_size + backend.tp_key_dim = backend.key_dim // tp_world_size + backend.tp_value_dim = backend.value_dim // tp_world_size + assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" + backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads + ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + backend.prefill_kernel = prefill_kernel + + +def _split_qkvzba(backend, mixed_qkvzba): + qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim + z_end = qkv_dim + backend.tp_value_dim + b_end = z_end + backend.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return mixed_qkv, z, b, a + + +def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + if decode: + batch_size = mixed_qkv.shape[0] + leading = (batch_size, 1) + else: + leading = (1, query.shape[0]) + query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) + return query, key, value + + +def _replay_state(att_state, layer_num: int) -> ReplayGDNState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=ReplayGDNState, + ) + + +def _single_replay_state(att_state) -> ReplayGDNState: + role = att_state.infer_state.att_state_role + matches = [ + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayGDNState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}") + return matches[0] diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 892c7901e3..65a8189e9d 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -40,27 +40,6 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo -def _replay_state(att_state, layer_num: int) -> ReplayKDAState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=ReplayKDAState, - ) - - -def _single_replay_state(att_state) -> ReplayKDAState: - role = att_state.infer_state.att_state_role - matches = [ - state - for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayKDAState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayKDAState for {role}") - return matches[0] - - class ReplayKDALinearAttBackend(BaseAttBackend): """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" @@ -162,37 +141,6 @@ def init_state(self): self.b_ssm_buffer_idx = self.b_conv_buffer_idx begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) - def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: - infer = self.infer_state - if mtp_step == 0: - self.b_conv_buffer_idx = infer.b_req_idx - self.b1_mtp_cu_q_seq_len = None - self.b_num_accepted_tokens = None - return - if self.backend.uses_dynamic_spec_verify_layout(): - cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=mtp_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - self.b1_mtp_cu_q_seq_len = cu - self.b_conv_buffer_idx = conv_idx - self.b_num_accepted_tokens = accepted - return - width = mtp_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "KDA replay decode requires batch_size divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - def decode_att( self, q: torch.Tensor, @@ -246,3 +194,55 @@ def decode_att( if self.b1_mtp_cu_q_seq_len is None: state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=True) return output + + def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if mtp_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = mtp_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "KDA replay decode requires batch_size divisible by mtp_step + 1, " + f"got batch_size={batch_size}, mtp_step={mtp_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + +def _replay_state(att_state, layer_num: int) -> ReplayKDAState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=ReplayKDAState, + ) + + +def _single_replay_state(att_state) -> ReplayKDAState: + role = att_state.infer_state.att_state_role + matches = [ + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayKDAState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}") + return matches[0] 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..266e2d2b49 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 @@ -3,50 +3,6 @@ import triton.language as tl -@triton.jit -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, - BLOCK_SIZE: tl.constexpr, -): - offsets = tl.arange(0, BLOCK_SIZE) - token_mask = offsets < batch_size - cu_mask = offsets <= batch_size - - 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) - - 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. - 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.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, @@ -109,3 +65,47 @@ def build_dynamic_mtp_linear_att_state_params( num_stages=1, ) return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + + +@triton.jit +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, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.arange(0, BLOCK_SIZE) + token_mask = offsets < batch_size + cu_mask = offsets <= batch_size + + 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) + + 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. + 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.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) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py index d1e1040213..5244dada83 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -5,28 +5,6 @@ import triton.language as tl -@triton.jit(do_not_specialize=["N"]) -def _commit_prefix( - Indices, - Reqs, - Lengths, - Cursors, - N, - HOLD: tl.constexpr, - REQ_STRIDE: tl.constexpr, - LENGTH_STRIDE: tl.constexpr, - BLOCK: tl.constexpr, -): - row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) - length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) - tl.store(Indices + req, length - 1, row < N) - if Cursors is not None: - valid = (row < N) & (req != HOLD) - cursor = tl.load(Cursors + req, valid, 0) - tl.store(Cursors + req, cursor + length, valid) - - def commit_ssm_prefix( indices: torch.Tensor, reqs: torch.Tensor, @@ -53,3 +31,25 @@ def commit_ssm_prefix( 256, num_warps=4, ) + + +@triton.jit(do_not_specialize=["N"]) +def _commit_prefix( + Indices, + Reqs, + Lengths, + Cursors, + N, + HOLD: tl.constexpr, + REQ_STRIDE: tl.constexpr, + LENGTH_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) + length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) + tl.store(Indices + req, length - 1, row < N) + if Cursors is not None: + valid = (row < N) & (req != HOLD) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor + length, valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index a3484c0d8f..72c5611704 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -174,16 +174,6 @@ def select_config( return run_config -def _replay_config(config: dict | None) -> dict[str, int]: - if not config: - return dict(_DEFAULT_REPLAY_CONFIG) - return { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } - - def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: """返回当前 batch 的固定或已调优 SSM kernel 配置。""" if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: @@ -199,3 +189,13 @@ def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens ssm_backend.batch_configs[key] = config logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") return config + + +def _replay_config(config: dict | None) -> dict[str, int]: + if not config: + return dict(_DEFAULT_REPLAY_CONFIG) + return { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index f1d0486966..60dfd016c8 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -117,37 +117,9 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config["autotune_layer_num"] = 4 return config - def _init_config(self): - with open(os.path.join(self.weight_dir_, "config.json")) as f: - self.config = self.normalize_model_config(json.load(f)) - - def _verify_params(self): - super()._verify_params() - assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" - assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" - if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): - raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") - if self.args.llm_kv_type == "fp8kv_dsa": - # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 - if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: - raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") - def autotune_layers(self): return 4 - def _init_some_value(self): - self.layers_num = self.config["n_layer"] - self.vocab_size = self.config["vocab_size"] - # MLA stores one replicated latent KV vector per token on every TP rank. - self.tp_k_head_num_ = 1 - self.tp_v_head_num_ = 0 - self.qk_nope_head_dim = self.config["qk_nope_head_dim"] - self.qk_rope_head_dim = self.config["qk_rope_head_dim"] - self.q_lora_rank = self.config["q_lora_rank"] - self.kv_lora_rank = self.config["kv_lora_rank"] - self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) - self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - @classmethod def create_full_att_states( cls, @@ -194,6 +166,34 @@ def prepare_attention_states(self) -> None: # Reserve fixed scratch space before the shared KV capacity is profiled. Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json")) as f: + self.config = self.normalize_model_config(json.load(f)) + + def _verify_params(self): + super()._verify_params() + assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" + assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): + raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa": + # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 + if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: + raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") + + def _init_some_value(self): + self.layers_num = self.config["n_layer"] + self.vocab_size = self.config["vocab_size"] + # MLA stores one replicated latent KV vector per token on every TP rank. + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + self.kv_lora_rank = self.config["kv_lora_rank"] + self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) + self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim + def _init_att_backend(self): from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( Fp8Glm5NextSparseAttBackend, diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 20fbeb4749..9da87c6acb 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -86,23 +86,9 @@ def create_bounded_att_states( ) return [state] - def _init_triton(self): - def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: - return torch.empty(size, device="cuda", dtype=torch.int8) - - # Set Triton allocator for TMA descriptors - # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py - triton.set_allocator(_triton_allocator) - logger.info("Triton allocator set for Qwen3Next model") - return - def autotune_layers(self): return self.config["full_attention_interval"] - def _init_config(self): - super()._init_config() - self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - @classmethod def create_full_att_states( cls, @@ -137,6 +123,20 @@ def create_full_att_states( ) return [state] + def _init_triton(self): + def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: + return torch.empty(size, device="cuda", dtype=torch.int8) + + # Set Triton allocator for TMA descriptors + # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py + triton.set_allocator(_triton_allocator) + logger.info("Triton allocator set for Qwen3Next model") + return + + def _init_config(self): + super()._init_config() + self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) + def _init_att_backend1(self): role = self.att_state_role if not self.att_state_manager.get_bounded_attention_states(role): From 5e7a2857704306ded21fb67162158a847bb35576 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 06:58:19 +0000 Subject: [PATCH 17/54] refactor: tune replay on a one-layer scratch state Autotune no longer wraps an existing SSM through a public constructor. The replay kernel entry only consumes that layer's tensors. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 78 +++++++++---------- .../bounded_att/kda/replay_kda_state.py | 78 +++++++++---------- .../attention/linear/replay_gdn_decode.py | 2 +- .../attention/linear/replay_kda_decode.py | 2 +- .../triton_kernel/linear_att/ssm_autotune.py | 28 ++++--- test/benchmark/kernels/benchmark_kda_ssm.py | 11 ++- .../linear_att/test_acceptance_fusion.py | 11 ++- .../linear_att/test_replayssm.py | 11 ++- .../linear_att/test_replayssm_execution.py | 11 ++- .../linear_att/test_ssm_autotune_memory.py | 11 ++- 10 files changed, 120 insertions(+), 123 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 0390a85e49..b44f4b57f2 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -100,47 +100,6 @@ def __init__( None, ) - @classmethod - def from_ssm( - cls, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - num_key_heads: int | None = None, - kda: bool = False, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: dict[str, int] | None = None, - ) -> "ReplayGDNState": - """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" - - if kda: - raise ValueError("GDN replay does not use the KDA gate") - owner = cls.__new__(cls) - owner.config = None - owner.state_name = None - owner.request_capacity = ssm.shape[1] - owner.mtp_step = verify_width - 1 - owner._algorithm = "gdn" - owner._logical_to_storage_layer = {} - owner.conv_state = None - owner.mtp_state_index = None - owner.big_pages_pool = None - owner.small_pages_pool = None - _bind_history( - owner, - ssm, - capacity, - verify_width, - activation_dtype, - num_key_heads, - gate_lower_bound, - compute_mode, - run_config, - ) - return owner - def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" return ( @@ -811,6 +770,43 @@ def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.T return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayGDNState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayGDNState.__new__(ReplayGDNState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + def _bind_history( owner, ssm: torch.Tensor, diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 5ac70f6818..9b53c01b01 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -100,47 +100,6 @@ def __init__( None, ) - @classmethod - def from_ssm( - cls, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - num_key_heads: int | None = None, - kda: bool = True, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: dict[str, int] | None = None, - ) -> "ReplayKDAState": - """用已有 SSM 张量组装 history。kernel 测试和 autotune scratch 走这里。""" - - if not kda: - raise ValueError("KDA replay requires the KDA gate") - owner = cls.__new__(cls) - owner.config = None - owner.state_name = None - owner.request_capacity = ssm.shape[1] - owner.mtp_step = verify_width - 1 - owner._algorithm = "kda" - owner._logical_to_storage_layer = {} - owner.conv_state = None - owner.mtp_state_index = None - owner.big_pages_pool = None - owner.small_pages_pool = None - _bind_history( - owner, - ssm, - capacity, - verify_width, - activation_dtype, - num_key_heads, - gate_lower_bound, - compute_mode, - run_config, - ) - return owner - def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" return ( @@ -810,6 +769,43 @@ def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.T return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayKDAState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayKDAState.__new__(ReplayKDAState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "kda" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + def _bind_history( owner, ssm: torch.Tensor, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 5844850c8e..27290c2f05 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -50,7 +50,7 @@ def replay( cu_seqlens: torch.Tensor | None = None, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 GDN replay decode，并返回输出与本轮 run config。""" + """执行一层 GDN replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" if history.kda: raise ValueError("GDN replay step does not use the KDA gate") diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index cfa59e10ab..4ee90de44a 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -50,7 +50,7 @@ def replay( cu_seqlens: torch.Tensor | None = None, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 KDA replay decode，并返回输出与本轮 run config。""" + """执行一层 KDA replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" if not history.kda: raise ValueError("KDA replay step requires the KDA gate") diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 72c5611704..60c67e50e4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -89,17 +89,27 @@ def rebuild_inputs( f"SSM tuning scratch needs {required} bytes, exceeding its {available // 2}-byte budget" ) _, _, hv, kd, vd = ssm_backend.ssm_state.shape - state = torch.empty((1, batch + 1, hv, kd, vd), device=q.device, dtype=ssm_backend.ssm_state.dtype).normal_(0, 0.02) - scratch = type(ssm_backend).from_ssm( - state, - ssm_backend.capacity, - ssm_backend.verify_width, - q.dtype, - num_key_heads=ssm_backend.num_key_heads, - kda=ssm_backend.kda, + from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig + + # 计时只需要一层、本 batch 的槽位。正式执行仍用调用方原来的 state。 + scratch = type(ssm_backend)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=ssm_backend.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=ssm_backend.ssm_state.dtype, + ), + batch + 1, + mtp_step=ssm_backend.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": ssm_backend.capacity, "compute_mode": ssm_backend.compute_mode}, gate_lower_bound=ssm_backend.gate_lower_bound, - compute_mode=ssm_backend.compute_mode, ) + scratch.ssm_state.normal_(0, 0.02) # benchmark 回调不能再次触发配置选择。 scratch._config_is_fixed = True inputs = [ diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index ea9020dd18..f4928a9017 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -17,14 +17,13 @@ def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.get("kda", False) + kda = kwargs.pop("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - return ReplayKDAState.from_ssm(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - return ReplayGDNState.from_ssm(*args, **kwargs) + return _history_owner(*args, **kwargs) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer 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 index fb5e9221e0..fb5168f661 100644 --- 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 @@ -6,14 +6,13 @@ def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.get("kda", False) + kda = kwargs.pop("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - return ReplayKDAState.from_ssm(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - return ReplayGDNState.from_ssm(*args, **kwargs) + return _history_owner(*args, **kwargs) def ssm_snapshot(backend, req): 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 index 94e583b435..34b16eec1c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -3,14 +3,13 @@ def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.get("kda", False) + kda = kwargs.pop("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - - return ReplayKDAState.from_ssm(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - return ReplayGDNState.from_ssm(*args, **kwargs) + return _history_owner(*args, **kwargs) def ssm_snapshot(backend, req): 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 index 7ed9dc7e90..33180c8285 100644 --- 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 @@ -4,14 +4,13 @@ def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.get("kda", False) + kda = kwargs.pop("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - return ReplayKDAState.from_ssm(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - return ReplayGDNState.from_ssm(*args, **kwargs) + return _history_owner(*args, **kwargs) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer 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 index 2627540fb2..9312f663c8 100644 --- 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 @@ -5,14 +5,13 @@ def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.get("kda", False) + kda = kwargs.pop("kda", False) if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - return ReplayKDAState.from_ssm(*args, **kwargs) - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - return ReplayGDNState.from_ssm(*args, **kwargs) + return _history_owner(*args, **kwargs) @pytest.mark.parametrize("width", [1, 5]) From 0e0884339e560342ddc2965e81d4d38fbe72e572 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 07:05:15 +0000 Subject: [PATCH 18/54] refactor: drop replay decode wrappers and restore native backends Native GDN and KDA did not need a method reorder. Replay decode now calls prepare and accept directly, and config selection stays in the autotuner. Co-authored-by: Cursor --- .../common/basemodel/attention/linear/gdn.py | 111 +++++++++--------- .../common/basemodel/attention/linear/kda.py | 96 +++++++-------- .../basemodel/attention/linear/replay_gdn.py | 6 +- .../attention/linear/replay_gdn_decode.py | 24 ---- .../basemodel/attention/linear/replay_kda.py | 6 +- .../attention/linear/replay_kda_decode.py | 24 ---- .../common/test_linear_state_contract.py | 4 +- 7 files changed, 115 insertions(+), 156 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 95b7fc157b..986214568f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -36,13 +36,8 @@ def __init__(self, model: "TpPartBaseModel"): def get_prefill_kernel(self): pass - 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_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - def _init_linear_layer_metadata(self, network_config, tp_world_size): + self.mtp_step = get_env_start_args().mtp_step # Linear attention specific dimensions @@ -105,9 +100,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_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": + return LinearAttDecodeAttState(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 @@ -205,6 +207,7 @@ def _gdn_prefill_kernel( @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 @@ -230,6 +233,54 @@ def init_state(self): 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 decode_att( self, q: torch.Tensor, @@ -281,54 +332,6 @@ def decode_att( ) return core_attn_out, z - 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 _gdn_decode_kernel( self, mixed_qkv: torch.Tensor, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 6d5ddc0cb3..cd2f8afb39 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -157,6 +157,54 @@ def init_state(self): 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] + def decode_att( self, q: torch.Tensor, @@ -214,51 +262,3 @@ def decode_att( num_accepted_tokens=self.b_num_accepted_tokens, ) return output - - 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] diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 8494b39f1b..6039d205e5 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -38,7 +38,7 @@ ) from lightllm.utils.envs_utils import get_env_start_args -from .replay_gdn_decode import begin_decode, replay +from .replay_gdn_decode import replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -159,7 +159,9 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) self._init_indices(draft_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + state.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + if self.b1_mtp_cu_q_seq_len is None: + state.accept_single_token(self.b_conv_buffer_idx) def decode_att( self, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 27290c2f05..c2695f0258 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -12,30 +12,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config -def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" - - history.prepare_decode(reqs, cu_seqlens) - if cu_seqlens is None: - history.accept_single_token(reqs) - - -def run_config( - history, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, -) -> dict[str, int]: - """按当前 kernel 形状选择或返回固定的 GDN replay 调优配置。""" - - return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def replay( history, linear_layer_idx: int, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 65a8189e9d..6229ceb66a 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -34,7 +34,7 @@ build_dynamic_mtp_linear_att_state_params, ) -from .replay_kda_decode import begin_decode, replay +from .replay_kda_decode import replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -139,7 +139,9 @@ def init_state(self): state = _single_replay_state(self) self._init_indices(self.backend.mtp_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - begin_decode(state, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + state.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + if self.b1_mtp_cu_q_seq_len is None: + state.accept_single_token(self.b_conv_buffer_idx) def decode_att( self, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 4ee90de44a..eb48c8df20 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -12,30 +12,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config -def begin_decode(history, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None) -> None: - """折叠即将溢出的 history。普通 decode 同时接受本轮唯一 token。""" - - history.prepare_decode(reqs, cu_seqlens) - if cu_seqlens is None: - history.accept_single_token(reqs) - - -def run_config( - history, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, -) -> dict[str, int]: - """按当前 kernel 形状选择或返回固定的 KDA replay 调优配置。""" - - return resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - - def replay( history, linear_layer_idx: int, diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 9b4a8b031b..feff2d1f08 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -4,7 +4,7 @@ from lightllm.common.attention_state.bounded_att import KDAState from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.linear.replay_gdn_decode import run_config as gdn_run_config +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -75,7 +75,7 @@ def test_replay_backend_publishes_detached_run_config(): ssm_backend.run_config = tuning ssm_backend._config_is_fixed = True ssm_backend.compute_mode = "inline" - config = gdn_run_config(ssm_backend, *([None] * 7)) + config = resolve_run_config(ssm_backend, "replay", *([None] * 7), None) tuning["BV"] = 128 assert config == {"BV": 32, "num_warps": 2, "num_stages": 3} From 476cfdcf4cc1d26565b60bac2743f9292aeb7e7d Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 09:05:58 +0000 Subject: [PATCH 19/54] refactor: keep replay history as one accepted prefix Decode appends into a single cursor-length buffer and compresses the whole prefix only when the next verify would exceed capacity, unless force_compress is set. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 137 +++++++++--------- .../bounded_att/kda/replay_kda_state.py | 137 +++++++++--------- .../basemodel/attention/linear/replay_gdn.py | 6 +- .../attention/linear/replay_gdn_decode.py | 2 +- .../basemodel/attention/linear/replay_kda.py | 6 +- .../attention/linear/replay_kda_decode.py | 2 +- .../triton_kernel/linear_att/replayssm.py | 5 +- .../linear_att/replayssm_compute.py | 87 ++++------- .../linear_att/replayssm_fold.py | 126 +--------------- .../linear_att/replayssm_materialize.py | 14 +- .../triton_kernel/linear_att/ssm_autotune.py | 4 +- .../linear_att/test_replayssm.py | 43 ++++-- .../linear_att/test_replayssm_execution.py | 54 ++++--- 13 files changed, 248 insertions(+), 375 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index b44f4b57f2..0855033e5f 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -16,9 +16,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _accept_updates, - _fold_history, _materialize, - _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -87,7 +85,7 @@ def __init__( 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 - capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) + capacity, compute_mode, force_compress = _replay_settings(ssm_replay_config, mtp_step) _bind_history( self, self.ssm_state, @@ -98,6 +96,7 @@ def __init__( gate_lower_bound, compute_mode, None, + force_compress, ) def get_state_type(self) -> str: @@ -208,12 +207,11 @@ def restore_bounded_att_runtime_from_pages( self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受前缀，折进 canonical SSM，并清空 history。""" + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - if self.ssm_state.is_cuda: - self.merge_accepted_updates(req_indices) + self.maybe_compress_prefix(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -406,16 +404,17 @@ def read_from_pd_page( # ======================================================================== def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + """最后一层收口。非 MTP 先记下这 1 个 token，再按容量决定是否整段压缩。Prefill 总是压缩。""" if self.config is None or not self.config.logical_layer_ids: return if logical_layer_id != self.config.logical_layer_ids[-1]: return if fold: - self.merge_accepted_updates(reqs) + self._note_new_token(reqs) + self.maybe_compress_prefix(reqs) else: - self.clear_history(reqs) + self.compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -427,55 +426,10 @@ def clear_history(self, req: int | torch.Tensor) -> None: self.cursors[req] = 0 def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" - positions = torch.empty_like(reqs) - layers, _, hv, kd, vd = self.ssm_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, - num_warps=4, - ) - programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) - _fold_history[(programs,)]( - self.ssm_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.ssm_state.stride(0), - self.raw_keys.stride(0), - self.raw_values.stride(0), - self.gates.stride(0), - self.betas.stride(0), - False, - num_warps=2, - num_stages=1, - ) - return positions + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: """更新 cursor，保留每个请求本轮已接受的前缀。""" @@ -496,13 +450,47 @@ def accept_single_token(self, reqs: torch.Tensor) -> None: self.accept_updates(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把已接受 history 折回 canonical SSM，并清空 cursor。""" + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" if reqs.numel() == 0: return self._materialize_accepted(reqs, self.ssm_state, snapshot=False) self.cursors.index_fill_(0, reqs.long(), 0) + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + used = self.cursors.index_select(0, reqs.long()) + overflow = (reqs != self.hold_req_id) & (used + self.verify_width > self.capacity) + if bool(overflow.any()): + self.compress_prefix(reqs[overflow]) + + def reserve_prefix(self, reqs: torch.Tensor, width: int | torch.Tensor) -> None: + """下一步写不下时，先把已接受前缀整段压进 SSM。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + used = self.cursors.index_select(0, reqs.long()) + need = width if isinstance(width, torch.Tensor) else used.new_full((), width) + overflow = (reqs != self.hold_req_id) & (used + need > self.capacity) + if bool(overflow.any()): + self.merge_accepted_updates(reqs[overflow]) + def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -594,6 +582,16 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== + def _note_new_token(self, reqs: torch.Tensor) -> None: + """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" + + if reqs.numel() == 0: + return + live = reqs[reqs != self.hold_req_id] + if live.numel() == 0: + return + self.cursors[live.long()] += 1 + def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 @@ -757,11 +755,12 @@ def _materialize_accepted( ) -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str, bool]: config = dict(config or {}) capacity = int(config.get("cache_len", 16)) compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode + force_compress = bool(config.get("force_compress", False)) + return capacity, compute_mode, force_compress def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: @@ -817,8 +816,9 @@ def _bind_history( gate_lower_bound: float, compute_mode: str, run_config: dict[str, int] | None, + force_compress: bool = False, ) -> None: - """在 owner 上创建 GDN replay 的双半区 history 和 cursor。""" + """在 owner 上创建 GDN replay 的连续前缀 history 和 cursor。""" assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 @@ -826,9 +826,9 @@ def _bind_history( owner.ssm_state = ssm owner.kda = False owner.gate_lower_bound = gate_lower_bound - owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 owner.capacity = capacity owner.verify_width = verify_width + owner.force_compress = bool(force_compress) owner.compute_mode = compute_mode owner.run_config = run_config owner._config_is_fixed = run_config is not None @@ -838,20 +838,19 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.forward_cursors = torch.zeros_like(owner.cursors) - # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 FP32。 + # 一段连续前缀，长度不超过 capacity。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + key_shape = (layers, slots, owner.num_key_heads, capacity, k) owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) class _AcceptedTensors(NamedTuple): diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 9b53c01b01..f466977a1f 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -16,9 +16,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _accept_updates, - _fold_history, _materialize, - _prepare_fold, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -87,7 +85,7 @@ def __init__( 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 - capacity, compute_mode = _replay_settings(ssm_replay_config, mtp_step) + capacity, compute_mode, force_compress = _replay_settings(ssm_replay_config, mtp_step) _bind_history( self, self.ssm_state, @@ -98,6 +96,7 @@ def __init__( gate_lower_bound, compute_mode, None, + force_compress, ) def get_state_type(self) -> str: @@ -208,12 +207,11 @@ def restore_bounded_att_runtime_from_pages( self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受前缀，折进 canonical SSM，并清空 history。""" + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - if self.ssm_state.is_cuda: - self.merge_accepted_updates(req_indices) + self.maybe_compress_prefix(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -405,16 +403,17 @@ def read_from_pd_page( # 六、Replay history # ======================================================================== def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """本状态最后一层结束时收口 history。fold 时先把已接受前缀折进 canonical SSM。""" + """最后一层收口。非 MTP 先记下这 1 个 token，再按容量决定是否整段压缩。Prefill 总是压缩。""" if self.config is None or not self.config.logical_layer_ids: return if logical_layer_id != self.config.logical_layer_ids[-1]: return if fold: - self.merge_accepted_updates(reqs) + self._note_new_token(reqs) + self.maybe_compress_prefix(reqs) else: - self.clear_history(reqs) + self.compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -426,55 +425,10 @@ def clear_history(self, req: int | torch.Tensor) -> None: self.cursors[req] = 0 def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """保存本轮 decode 的 history 起点，并折叠即将溢出的旧半区。""" + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" - positions = torch.empty_like(reqs) - layers, _, hv, kd, vd = self.ssm_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, - num_warps=4, - ) - programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) - _fold_history[(programs,)]( - self.ssm_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.ssm_state.stride(0), - self.raw_keys.stride(0), - self.raw_values.stride(0), - self.gates.stride(0), - self.betas.stride(0), - True, - num_warps=2, - num_stages=1, - ) - return positions + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: """更新 cursor，保留每个请求本轮已接受的前缀。""" @@ -495,13 +449,47 @@ def accept_single_token(self, reqs: torch.Tensor) -> None: self.accept_updates(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把已接受 history 折回 canonical SSM，并清空 cursor。""" + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" if reqs.numel() == 0: return self._materialize_accepted(reqs, self.ssm_state, snapshot=False) self.cursors.index_fill_(0, reqs.long(), 0) + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + used = self.cursors.index_select(0, reqs.long()) + overflow = (reqs != self.hold_req_id) & (used + self.verify_width > self.capacity) + if bool(overflow.any()): + self.compress_prefix(reqs[overflow]) + + def reserve_prefix(self, reqs: torch.Tensor, width: int | torch.Tensor) -> None: + """下一步写不下时，先把已接受前缀整段压进 SSM。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + used = self.cursors.index_select(0, reqs.long()) + need = width if isinstance(width, torch.Tensor) else used.new_full((), width) + overflow = (reqs != self.hold_req_id) & (used + need > self.capacity) + if bool(overflow.any()): + self.merge_accepted_updates(reqs[overflow]) + def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -593,6 +581,16 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== + def _note_new_token(self, reqs: torch.Tensor) -> None: + """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" + + if reqs.numel() == 0: + return + live = reqs[reqs != self.hold_req_id] + if live.numel() == 0: + return + self.cursors[live.long()] += 1 + def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 @@ -756,11 +754,12 @@ def _materialize_accepted( ) -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str]: +def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str, bool]: config = dict(config or {}) capacity = int(config.get("cache_len", 16)) compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - return capacity, compute_mode + force_compress = bool(config.get("force_compress", False)) + return capacity, compute_mode, force_compress def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: @@ -816,8 +815,9 @@ def _bind_history( gate_lower_bound: float, compute_mode: str, run_config: dict[str, int] | None, + force_compress: bool = False, ) -> None: - """在 owner 上创建 KDA replay 的双半区 history 和 cursor。""" + """在 owner 上创建 KDA replay 的连续前缀 history 和 cursor。""" assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 @@ -825,9 +825,9 @@ def _bind_history( owner.ssm_state = ssm owner.kda = True owner.gate_lower_bound = gate_lower_bound - owner.fold_programs = 8 * torch.cuda.get_device_properties(ssm.device).multi_processor_count if ssm.is_cuda else 0 owner.capacity = capacity owner.verify_width = verify_width + owner.force_compress = bool(force_compress) owner.compute_mode = compute_mode owner.run_config = run_config owner._config_is_fixed = run_config is not None @@ -837,20 +837,19 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.forward_cursors = torch.zeros_like(owner.cursors) - # fold 时交替使用两个半区。precompute 且 BF16、短 history 时 history 用激活精度，其余用 FP32。 + # 一段连续前缀，长度不超过 capacity。precompute 且 BF16、短 history 时 history 用激活精度，其余用 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, owner.num_key_heads, 2 * capacity, k) + key_shape = (layers, slots, owner.num_key_heads, capacity, k) owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=torch.float32, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=history_dtype, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=ssm.device) owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) class _AcceptedTensors(NamedTuple): diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 6039d205e5..85a0077380 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -159,9 +159,11 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) self._init_indices(draft_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - state.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) if self.b1_mtp_cu_q_seq_len is None: - state.accept_single_token(self.b_conv_buffer_idx) + state.reserve_prefix(self.b_conv_buffer_idx, 1) + else: + cu = self.b1_mtp_cu_q_seq_len + state.reserve_prefix(self.b_conv_buffer_idx, cu[1:] - cu[:-1]) def decode_att( self, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index c2695f0258..d17d8e06e7 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -55,7 +55,7 @@ def replay( history.raw_values[linear_layer_idx], history.betas[linear_layer_idx], reqs, - history.forward_cursors, + history.cursors, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 6229ceb66a..9d988a9840 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -139,9 +139,11 @@ def init_state(self): state = _single_replay_state(self) self._init_indices(self.backend.mtp_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - state.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) if self.b1_mtp_cu_q_seq_len is None: - state.accept_single_token(self.b_conv_buffer_idx) + state.reserve_prefix(self.b_conv_buffer_idx, 1) + else: + cu = self.b1_mtp_cu_q_seq_len + state.reserve_prefix(self.b_conv_buffer_idx, cu[1:] - cu[:-1]) def decode_att( self, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index eb48c8df20..c4208eae31 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -55,7 +55,7 @@ def replay( history.raw_values[linear_layer_idx], history.betas[linear_layer_idx], reqs, - history.forward_cursors, + history.cursors, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 81b2804132..a8751ea970 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -12,18 +12,15 @@ _replay_gdn_window, _replay_sequential, ) -from .replayssm_fold import _accept_updates, _fold_history, _prepare_fold, _use_separate_fold +from .replayssm_fold import _accept_updates from .replayssm_materialize import _materialize __all__ = [ "_accept_updates", - "_fold_history", "_materialize", - "_prepare_fold", "_project_replay_queries", "_replay", "_replay_gdn_window", "_replay_sequential", "_replay_update", - "_use_separate_fold", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py index 604c09f1e0..c40c5bf44f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py @@ -4,8 +4,6 @@ import triton import triton.language as tl -from .replayssm_common import _replay_update - @triton.jit def _project_replay_queries( @@ -26,7 +24,6 @@ def _project_replay_queries( KDA: tl.constexpr, LOWER: tl.constexpr, BQK: tl.constexpr, - base, end, hv, kk, @@ -59,7 +56,7 @@ def _project_replay_queries( if KDA: history = ll < n gates = tl.load( - Gates + (slot * (2 * L) + base + ll[:, None]) * K + kk[None, :], + Gates + (slot * L + ll[:, None]) * K + kk[None, :], history[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) @@ -117,7 +114,6 @@ def _replay_gdn_window( BV: tl.constexpr, BW: tl.constexpr, BR: tl.constexpr, - base, end, h, hv, @@ -143,7 +139,7 @@ def _replay_gdn_window( 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, :], + Keys + (key_slot * L) * K + jj[:, None] * K + kk[None, :], old[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) @@ -162,7 +158,7 @@ def _replay_gdn_window( 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.load(Gates + slot * L + 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) @@ -173,7 +169,7 @@ def _replay_gdn_window( 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, :], + Deltas + (slot * L) * V + jj[:, None] * V + vv[None, :], old[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) @@ -194,31 +190,32 @@ def _replay_gdn_window( outputs, (valid & (tt < BW))[:, None] & (vv[None, :] < V), ) - records = base + n + jj - L + records = n + jj - L + in_prefix = fresh & (records < L) tl.store( - Deltas + (slot * (2 * L) + records[:, None]) * V + vv[None, :], + Deltas + (slot * L + records[:, None]) * V + vv[None, :], deltas, - fresh[:, None] & (vv[None, :] < V), + in_prefix[:, None] & (vv[None, :] < V), ) tl.store( - RawValues + (slot * (2 * L) + records[:, None]) * V + vv[None, :], + RawValues + (slot * L + records[:, None]) * V + vv[None, :], values, - fresh[:, None] & (vv[None, :] < V), + in_prefix[:, None] & (vv[None, :] < V), ) if iv == 0: if hv % (HV // KH) == 0: tl.store( - Keys + (key_slot * (2 * L) + records[:, None]) * K + kk[None, :], + Keys + (key_slot * L + records[:, None]) * K + kk[None, :], fresh_keys, - fresh[:, None] & (kk[None, :] < K), + in_prefix[:, None] & (kk[None, :] < K), ) tl.store( - RawKeys + (key_slot * (2 * L) + records[:, None]) * K + kk[None, :], + RawKeys + (key_slot * L + records[:, None]) * K + kk[None, :], raw_keys, - fresh[:, None] & (kk[None, :] < K), + in_prefix[:, None] & (kk[None, :] < K), ) - tl.store(Gates + slot * (2 * L) + records, fresh_gates, fresh) - tl.store(Betas + slot * (2 * L) + records, betas, fresh) + tl.store(Gates + slot * L + records, fresh_gates, in_prefix) + tl.store(Betas + slot * L + records, betas, in_prefix) return @@ -253,7 +250,6 @@ def _replay_sequential( BW: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, - base, end, hv, iv, @@ -271,23 +267,23 @@ def _replay_sequential( # KDA、inline 以及不满足 precompute 条件的路径使用逐 token 递推。 history = ll < n keys = tl.load( - Keys + (key_slot * (2 * L) + base) * K + ll[:, None] * K + kk[None, :], + Keys + (key_slot * L) * 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, :], + Deltas + (slot * L) * V + ll[:, None] * V + vv[None, :], history[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) if KDA: gates = tl.load( - Gates + (slot * (2 * L) + base + ll[:, None]) * K + kk[None, :], + Gates + (slot * L + 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) + gates = tl.load(Gates + slot * L + ll, history, 0).to(tl.float32) gate_prefix = tl.cumsum(gates, axis=0) total_g = tl.sum(gates, axis=0) if KDA: @@ -334,20 +330,21 @@ def _replay_sequential( 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) + record = n + in_prefix = record < L + tl.store(Deltas + (slot * L + record) * V + vv, d, in_prefix & (vv < V)) + tl.store(RawValues + (slot * L + record) * V + vv, v, in_prefix & (vv < V)) if iv == 0: # Only the group owner writes shared K. Every CTA computes this # round's K locally; readers only consume earlier accepted records. if hv % (HV // KH) == 0: - tl.store(Keys + (key_slot * (2 * L) + record) * K + kk, k, kk < K) - tl.store(RawKeys + (key_slot * (2 * L) + record) * K + kk, raw_k, kk < K) + tl.store(Keys + (key_slot * L + record) * K + kk, k, in_prefix & (kk < K)) + tl.store(RawKeys + (key_slot * L + record) * K + kk, raw_k, in_prefix & (kk < K)) if KDA: - tl.store(Gates + (slot * (2 * L) + record) * K + kk, g, kk < K) + tl.store(Gates + (slot * L + record) * K + kk, g, in_prefix & (kk < K)) else: - tl.store(Gates + slot * (2 * L) + record, g) - tl.store(Betas + slot * (2 * L) + record, beta) + tl.store(Gates + slot * L + record, g, in_prefix) + tl.store(Betas + slot * L + record, beta, in_prefix) keys = tl.where((ll == n)[:, None], k[None, :], keys) deltas = tl.where((ll == n)[:, None], d[None, :], deltas) if not KDA: @@ -400,8 +397,8 @@ def _replay( ): """执行一个请求窗口的 replay decode，并写入 history。 - 当 history 即将跨半区时，先按 raw input 折叠到 canonical state；随后 - 根据 compute mode 选择预计算投影或逐 token 递推。请求数属于运行时布局， + 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 + cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， 不应因为 eager batch 改变而形成新的 kernel 变体。 """ iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -422,27 +419,10 @@ def _replay( 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 + n = tl.load(Cursors + req) 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 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 @@ -466,7 +446,6 @@ def _replay( KDA, LOWER, BQK, - base, end, hv, kk, @@ -508,7 +487,6 @@ def _replay( BV, BW, BR, - base, end, h, hv, @@ -558,7 +536,6 @@ def _replay( BW, KDA, LOWER, - base, end, hv, iv, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 3281c5c402..136ead33f8 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -"""ReplaySSM history preparation and accepted-prefix folding kernels.""" +"""ReplaySSM accepted-prefix length updates.""" import triton import triton.language as tl -from .replayssm_common import _replay_update - @triton.jit(do_not_specialize=["N"]) def _accept_updates( @@ -24,125 +22,3 @@ def _accept_updates( 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 _use_separate_fold(valid, length, count, N, L: tl.constexpr, WIDTH: tl.constexpr): - # 只决定本轮是否单独启动 fold；真正的 history 合并由 _fold_history 完成。 - small_batch: tl.constexpr = 12 - dense_batch: tl.constexpr = 256 - live = tl.sum((valid & (length > 0)).to(tl.int32), axis=0) - split = (count > 0) & ((live <= small_batch) | (2 * count <= live)) - if L == 8 and WIDTH == 4: - if N >= dense_batch: - short = tl.sum((valid & (length > 0) & (length < WIDTH)).to(tl.int32), axis=0) - split = split & ((live < dense_batch) | (short > 0)) - return split - - -@triton.jit(do_not_specialize=["N"]) -def _prepare_fold( - Reqs, - Cursors, - Positions, - ForwardCursors, - Cu, - Active, - N, - HOLD: tl.constexpr, - L: tl.constexpr, - VARLEN: tl.constexpr, - BLOCK: tl.constexpr, - WIDTH: tl.constexpr, -): - # 输出 Positions/ForwardCursors 供 decode 使用，Active/N 供后续 fold - # 使用；只改写本轮请求对应的 cursor，不触碰 canonical state。 - 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) - split = _use_separate_fold(valid, length, count, N, L, WIDTH) - 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(do_not_specialize=["N"]) -def _fold_history( - State, - RawKeys, - RawValues, - Gates, - Betas, - Reqs, - Cursors, - Active, - N, - 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, -): - # 只消费 Active 中的 fold 请求，把 history 合并回 canonical state； - # cursor 的 phase 已由 _prepare_fold 切换，写入后新 history 从同一半区继续。 - 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) - 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) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py index 3b19a13e6c..a49312e825 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py @@ -36,9 +36,7 @@ def _materialize( 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 + n = tl.load(Cursors + req) if n == 0 and not SNAPSHOT and not FLUSH_EXPORT: return linear_layer_idx, head = lh // HV, lh % HV @@ -49,13 +47,13 @@ def _materialize( 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) + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + 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) + g = tl.load(Gates + (slot * L + 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) + g = tl.load(Gates + slot * L + j).to(tl.float32) + beta = tl.load(Betas + slot * L + 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)) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 60c67e50e4..5aaef7a377 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -174,9 +174,9 @@ def select_config( ssm_backend.run_config = run_config reqs, accepted = workload ssm_backend.cursors.zero_() - # 覆盖完整 history 周期并触发一次 fold；只重复第一轮会一直测量空 history。 + # 把前缀写满但不越出 capacity。只重复空 history 会测不到已有前缀的读取。 tokens = q.shape[1 if cu_seqlens is not None else 0] - rounds = ssm_backend.capacity // min(ssm_backend.verify_width, tokens) + 1 + rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) for step in range(rounds): ssm_backend.prepare_decode(reqs, cu_seqlens) replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) 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 index 34b16eec1c..9aab192432 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -201,7 +201,9 @@ def step(): 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: + query_len = width if seq == 0 else max(width - 1, 1) + if pending[seq] + query_len > capacity: + ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) expected[:, req] = expected[:, req].to(dtype).float() pending[seq] = 0 pending[seq] += counts[seq] @@ -337,6 +339,19 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d accepted[0] = step % max(1, width - 1) accepted[1] = step % width outputs = [] + step_lengths = (cu[1:] - cu[:-1]).tolist() + if any( + ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + for req, length in zip(reqs.tolist(), step_lengths) + ): + overflow = [ + req + for req, length in zip(reqs.tolist(), step_lengths) + if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) for ssm_backend in ssm_backends: ssm_backend.prepare_decode(reqs, cu) outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) @@ -353,9 +368,8 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_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. +def test_bf16_commit_rounds_before_the_next_token(): + # 前缀写满后先压回 SSM，下一 token 从长度 0 重写。两次 commit 必须得到同一块 BF16 状态。 state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) ssm_backends = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] reqs = torch.tensor([0], device="cuda", dtype=torch.int32) @@ -365,7 +379,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): ssm_backend.gates.fill_(-0.01) ssm_backend.betas.zero_() ssm_backend.cursors[0] = 16 - ssm_backends[1].merge_accepted_updates(reqs) + ssm_backend.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") @@ -375,8 +389,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): bias = torch.zeros_like(alog) outputs = [] for ssm_backend in ssm_backends: - ssm_backend.prepare_decode(reqs) - ssm_backend.accept_updates(reqs) + assert ssm_backend.cursors[0].item() == 0 outputs.append(replay_layer(ssm_backend, 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].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) @@ -392,7 +405,7 @@ def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, 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] + pending = [0, 0] reqs = torch.full((3,), 3, dtype=torch.int32, device="cuda") cu = torch.tensor([0, 4, 8, 8], dtype=torch.int32, device="cuda") accepted = torch.zeros(4, dtype=torch.int32, device="cuda") @@ -436,12 +449,12 @@ def step(): 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: + ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) expected[:, req] = expected[:, req].to(dtype).float() pending[req] = 0 - phases[req] ^= 16 + old_cursors = [pending[req] for req in ids] + [0] if not capture_metadata: positions.copy_(ssm_backend.prepare_decode(reqs, cu)) graph.replay() @@ -458,12 +471,12 @@ def step(): 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] + assert ssm_backend.cursors.tolist() == pending + [0, 0] for output in outputs: assert torch.count_nonzero(output[:, sum(widths) :]).item() == 0 torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) if iteration == 2: - assert pending[0] == 8 and phases[0] == 0 + assert pending[0] == 8 torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) ssm_backend.merge_accepted_updates(reqs[:2]) torch.testing.assert_close( @@ -631,11 +644,9 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa k, v = torch.randn_like(q), torch.randn_like(q) a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") b = torch.randn(live, query_width, 1, device="cuda") + if (backend.cursors[:live] + query_width > backend.capacity).any(): + backend.merge_accepted_updates(reqs[:live]) backend.prepare_decode(reqs, cu) - if step == 2: - inline = backend.forward_cursors[:live] % 16 + query_width > 8 - assert inline[:folding].eq(not separate).all() - assert not inline[folding:].any() actual = replay_layer( backend, 0, 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 index 33180c8285..5ca0d1186b 100644 --- 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 @@ -1,6 +1,5 @@ import pytest import torch -import triton def ReplaySSMExecutor(*args, **kwargs): @@ -26,7 +25,7 @@ def ReplaySSMExecutor(*args, **kwargs): def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): from lightllm.common.basemodel.triton_kernel.linear_att import replayssm - kernels = ("_prepare_fold", "_fold_history", "_replay", "_accept_updates") + kernels = ("_replay", "_accept_updates") hashes = {name: set() for name in kernels} launches = [] for name in kernels: @@ -75,7 +74,11 @@ def cycle(backend, ids): graph, ids, output = graphs[32 if count <= 32 else 512] ids.fill_(captured.hold_req_id) ids[:count].copy_(reqs[:count]) - for _ in range(9): + for _ in range(8): + live = reqs[:count] + if (eager.cursors[live.long()] + 1 > eager.capacity).any(): + eager.merge_accepted_updates(live) + captured.merge_accepted_updates(ids[:count]) expected = cycle(eager, reqs[:count]) calls = len(launches) graph.replay() @@ -83,9 +86,8 @@ def cycle(backend, ids): torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) assert torch.equal(captured.cursors, eager.cursors) torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) - assert len(hashes["_prepare_fold"]) == len({triton.next_power_of_2(n) for n in counts}) - for name in kernels[1:]: - assert len(hashes[name]) == 1, (name, hashes[name]) + assert len(hashes["_replay"]) == 1 + assert len(hashes["_accept_updates"]) == 1 assert captured.cursors[-1] == 0 assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) @@ -192,6 +194,14 @@ def forward_accept(index, cu, reqs): ids, lengths, counts, logical = inputs(iteration) for tensor in [q, k, v, beta]: tensor.normal_() + overflow = [ + req for req, length in zip(ids, lengths) if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) + folds += len(overflow) for index, ssm_backend in enumerate(ssm_backends): if not capture_metadata: for dst, src in zip(captured[index], metadata(ssm_backend)): @@ -201,8 +211,6 @@ def forward_accept(index, cu, reqs): 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 @@ -261,7 +269,7 @@ def warm(ssm_backend): return out graphs, positions, outputs = [], [], [] - for early, ssm_backend in enumerate(ssm_backends): + for ssm_backend in ssm_backends: warm(ssm_backend) ssm_backend.ssm_state.copy_(initial) ssm_backend.cursors.zero_() @@ -275,27 +283,27 @@ def warm(ssm_backend): 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: + if capture_metadata: 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_() + overflow = [int(req) for req in reqs.tolist() if ssm_backends[0].cursors[req] + 1 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) old = ssm_backends[0].cursors[reqs.long()].clone() - for early, (ssm_backend, graph, pos) in enumerate(zip(ssm_backends, graphs, positions)): + for ssm_backend, graph, pos in 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: + if not capture_metadata: ssm_backend.accept_updates(reqs) torch.cuda.synchronize() assert torch.equal(positions[0], old) @@ -309,10 +317,6 @@ def warm(ssm_backend): 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(): @@ -340,6 +344,9 @@ def test_cross_layer_fold_with_state_offset_above_int32(): cu = reqs.new_tensor([0, 4, 8]) def step(ssm_backend, accepted): + overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] + if overflow: + ssm_backend.merge_accepted_updates(reqs.new_tensor(overflow)) pos = ssm_backend.prepare_decode(reqs, cu) out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] ssm_backend.accept_updates(reqs, accepted) @@ -408,6 +415,11 @@ def step(): accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 for value in [q, k, v, beta]: value.normal_() + overflow = [idx for idx in range(batch) if ssm_backends[0].cursors[idx] + 4 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(overflow_reqs) for graph in graphs: graph.replay() for small, padded_out in zip(outputs[0], outputs[1]): From 837406404f5f64af78c31ac913673d4d2082b602 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 09:14:56 +0000 Subject: [PATCH 20/54] refactor: compress replay prefixes inside the materialize kernel Decide whether the accepted prefix fits the next verify on device, so CUDA graph replay keeps the same rule without a host-side branch. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 19 ++++------------ .../bounded_att/kda/replay_kda_state.py | 19 ++++------------ .../basemodel/attention/linear/replay_gdn.py | 5 ----- .../basemodel/attention/linear/replay_kda.py | 5 ----- .../linear_att/replayssm_materialize.py | 22 ++++++++++++------- 5 files changed, 22 insertions(+), 48 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 0855033e5f..9820efbe69 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -475,21 +475,7 @@ def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: if self.force_compress: self.compress_prefix(reqs) return - used = self.cursors.index_select(0, reqs.long()) - overflow = (reqs != self.hold_req_id) & (used + self.verify_width > self.capacity) - if bool(overflow.any()): - self.compress_prefix(reqs[overflow]) - - def reserve_prefix(self, reqs: torch.Tensor, width: int | torch.Tensor) -> None: - """下一步写不下时，先把已接受前缀整段压进 SSM。""" - - if reqs.numel() == 0 or not self.ssm_state.is_cuda: - return - used = self.cursors.index_select(0, reqs.long()) - need = width if isinstance(width, torch.Tensor) else used.new_full((), width) - overflow = (reqs != self.hold_req_id) & (used + need > self.capacity) - if bool(overflow.any()): - self.merge_accepted_updates(reqs[overflow]) + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -721,6 +707,7 @@ def _materialize_accepted( snapshot: bool, output_strides=None, flush_export: bool = False, + when_full: bool = False, ) -> None: """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" @@ -750,6 +737,8 @@ def _materialize_accepted( False, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, num_warps=config["num_warps"], num_stages=config["num_stages"], ) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index f466977a1f..c56b13cb6d 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -474,21 +474,7 @@ def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: if self.force_compress: self.compress_prefix(reqs) return - used = self.cursors.index_select(0, reqs.long()) - overflow = (reqs != self.hold_req_id) & (used + self.verify_width > self.capacity) - if bool(overflow.any()): - self.compress_prefix(reqs[overflow]) - - def reserve_prefix(self, reqs: torch.Tensor, width: int | torch.Tensor) -> None: - """下一步写不下时，先把已接受前缀整段压进 SSM。""" - - if reqs.numel() == 0 or not self.ssm_state.is_cuda: - return - used = self.cursors.index_select(0, reqs.long()) - need = width if isinstance(width, torch.Tensor) else used.new_full((), width) - overflow = (reqs != self.hold_req_id) & (used + need > self.capacity) - if bool(overflow.any()): - self.merge_accepted_updates(reqs[overflow]) + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -720,6 +706,7 @@ def _materialize_accepted( snapshot: bool, output_strides=None, flush_export: bool = False, + when_full: bool = False, ) -> None: """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" @@ -749,6 +736,8 @@ def _materialize_accepted( True, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, num_warps=config["num_warps"], num_stages=config["num_stages"], ) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 85a0077380..c15c5bf694 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -159,11 +159,6 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) self._init_indices(draft_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - if self.b1_mtp_cu_q_seq_len is None: - state.reserve_prefix(self.b_conv_buffer_idx, 1) - else: - cu = self.b1_mtp_cu_q_seq_len - state.reserve_prefix(self.b_conv_buffer_idx, cu[1:] - cu[:-1]) def decode_att( self, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 9d988a9840..dd8650fbcb 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -139,11 +139,6 @@ def init_state(self): state = _single_replay_state(self) self._init_indices(self.backend.mtp_step, state.mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx - if self.b1_mtp_cu_q_seq_len is None: - state.reserve_prefix(self.b_conv_buffer_idx, 1) - else: - cu = self.b1_mtp_cu_q_seq_len - state.reserve_prefix(self.b_conv_buffer_idx, cu[1:] - cu[:-1]) def decode_att( self, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py index a49312e825..acf50cb5c7 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py @@ -4,6 +4,8 @@ import triton import triton.language as tl +from .replayssm_common import _replay_update + @triton.jit def _materialize( @@ -28,16 +30,21 @@ def _materialize( KDA: tl.constexpr, FLUSH_EXPORT: tl.constexpr = False, OUTPUT_STRIDES: tl.constexpr = None, + WHEN_FULL: tl.constexpr = False, + NEXT: tl.constexpr = 1, ): - # 将 cursor 指向的 accepted history 应用到 canonical state；SNAPSHOT - # 只写独立输出，merge 路径则原地更新请求槽位。FLUSH_EXPORT 同时 - # 更新 canonical state 和缓存目标，并在同一个 kernel 中清空 cursor。 + # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT + # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， + # 并在 kernel 内把这些请求的 cursor 清零，避免主机分支进不了 CUDA Graph。 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 n = tl.load(Cursors + req) - if n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + if WHEN_FULL: + if n == 0 or n + NEXT <= L: + return + elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: return linear_layer_idx, head = lh // HV, lh % HV slot = (linear_layer_idx * SLOTS + req) * HV + head @@ -54,10 +61,7 @@ def _materialize( else: g = tl.load(Gates + slot * L + j).to(tl.float32) beta = tl.load(Betas + slot * L + 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, :] + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) if FLUSH_EXPORT: target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) target_out = ( @@ -85,3 +89,5 @@ def _materialize( 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)) + if WHEN_FULL and iv == 0 and lh == 0: + tl.store(Cursors + req, 0) From 0ead21cc78f6fad53c0a3325ec87b9313bce312b Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 10:28:23 +0000 Subject: [PATCH 21/54] refactor: close replay tokens in-kernel and drop tuning OOM bypass Record each accepted decode token with the cursor kernel so CUDA graph capture does not synchronize on request indices. Tune the update on a one-layer scratch, and let a real scratch allocation failure surface instead of being remembered as an OOM key. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 44 +++++++++----- .../bounded_att/kda/replay_kda_state.py | 44 +++++++++----- .../basemodel/attention/linear/replay_gdn.py | 11 ++-- .../basemodel/attention/linear/replay_kda.py | 5 +- .../triton_kernel/linear_att/ssm_autotune.py | 39 ++----------- lightllm/common/triton_utils/autotuner.py | 21 +------ .../linear_att/test_ssm_autotune_memory.py | 48 ++++++++------- unit_tests/common/test_autotuner.py | 58 ------------------- 8 files changed, 96 insertions(+), 174 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 9820efbe69..6bb90bf064 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -403,19 +403,20 @@ def read_from_pd_page( # 六、Replay history # ======================================================================== - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """最后一层收口。非 MTP 先记下这 1 个 token，再按容量决定是否整段压缩。Prefill 总是压缩。""" + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" - if self.config is None or not self.config.logical_layer_ids: - return - if logical_layer_id != self.config.logical_layer_ids[-1]: - return - if fold: - self._note_new_token(reqs) - self.maybe_compress_prefix(reqs) - else: + if self._is_last_layer(logical_layer_id): self.compress_prefix(reqs) + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + if not self._is_last_layer(logical_layer_id): + return + self._note_new_token(reqs) + self.maybe_compress_prefix(reqs) + def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -445,9 +446,9 @@ def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tenso ) def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 把单 token 记入当前 history。""" + """把一个已接受 token 记入前缀长度。""" - self.accept_updates(reqs) + self._note_new_token(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -568,15 +569,26 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== + def _is_last_layer(self, logical_layer_id: int) -> bool: + if self.config is None or not self.config.logical_layer_ids: + return False + return logical_layer_id == self.config.logical_layer_ids[-1] + def _note_new_token(self, reqs: torch.Tensor) -> None: """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" if reqs.numel() == 0: return - live = reqs[reqs != self.hold_req_id] - if live.numel() == 0: - return - self.cursors[live.long()] += 1 + # 固定加 1，并在 kernel 内跳过 hold。避免按请求号做布尔索引，CUDA graph 不能捕获这种同步。 + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs, + reqs.numel(), + self.hold_req_id, + 1, + 256, + ) def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index c56b13cb6d..73b6253452 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -402,19 +402,20 @@ def read_from_pd_page( # ======================================================================== # 六、Replay history # ======================================================================== - def close_history_after_layer(self, logical_layer_id: int, reqs: torch.Tensor, fold: bool) -> None: - """最后一层收口。非 MTP 先记下这 1 个 token，再按容量决定是否整段压缩。Prefill 总是压缩。""" + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" - if self.config is None or not self.config.logical_layer_ids: - return - if logical_layer_id != self.config.logical_layer_ids[-1]: - return - if fold: - self._note_new_token(reqs) - self.maybe_compress_prefix(reqs) - else: + if self._is_last_layer(logical_layer_id): self.compress_prefix(reqs) + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + if not self._is_last_layer(logical_layer_id): + return + self._note_new_token(reqs) + self.maybe_compress_prefix(reqs) + def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -444,9 +445,9 @@ def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tenso ) def accept_single_token(self, reqs: torch.Tensor) -> None: - """非 MTP decode 把单 token 记入当前 history。""" + """把一个已接受 token 记入前缀长度。""" - self.accept_updates(reqs) + self._note_new_token(reqs) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -567,15 +568,26 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== + def _is_last_layer(self, logical_layer_id: int) -> bool: + if self.config is None or not self.config.logical_layer_ids: + return False + return logical_layer_id == self.config.logical_layer_ids[-1] + def _note_new_token(self, reqs: torch.Tensor) -> None: """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" if reqs.numel() == 0: return - live = reqs[reqs != self.hold_req_id] - if live.numel() == 0: - return - self.cursors[live.long()] += 1 + # 固定加 1，并在 kernel 内跳过 hold。避免按请求号做布尔索引，CUDA graph 不能捕获这种同步。 + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs, + reqs.numel(), + self.hold_req_id, + 1, + 256, + ) def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index c15c5bf694..eb6b647e1f 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -91,7 +91,6 @@ class ReplayGDNPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - # Prefill 直接覆盖 canonical SSM。进入前 history 已经在 decode 接受时折完。 reqs = self.infer_state.b_req_idx self.b_conv_buffer_idx = reqs self.b_ssm_buffer_idx = reqs @@ -115,7 +114,7 @@ def prefill_att( conv_states = conv_states[:, :, : -backend.mtp_step] mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) - state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=False) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) return core_attn_out, z def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): @@ -179,8 +178,11 @@ def decode_att( draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) if draft_step > 0: core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) - return core, z - return self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + else: + core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + if self.b1_mtp_cu_q_seq_len is None: + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + return core, z def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: infer = self.infer_state @@ -244,7 +246,6 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) - state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=True) return output, z def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index dd8650fbcb..ece6f309b5 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -71,7 +71,6 @@ class ReplayKDAPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - # Prefill 直接覆盖 canonical SSM。进入前 history 已经在 decode 接受时折完。 reqs = self.infer_state.b_req_idx self.b_conv_buffer_idx = reqs self.b_ssm_buffer_idx = reqs @@ -124,7 +123,7 @@ def prefill_att( lower_bound=backend.lower_bound, ) ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=False) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) return output @@ -191,7 +190,7 @@ def decode_att( self.b1_mtp_cu_q_seq_len, ) if self.b1_mtp_cu_q_seq_len is None: - state.close_history_after_layer(layer_num, self.b_conv_buffer_idx, fold=True) + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) return output def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 5aaef7a377..44d8d7178d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -26,15 +26,13 @@ def get_configs() -> list[dict[str, int]]: def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens) -> dict[str, object]: """构造区分状态布局、数据类型和投影模式的静态调优键。""" - axis = 1 if cu_seqlens is not None else 0 _, _, hv, kd, vd = ssm_backend.ssm_state.shape # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 key = { "mode": mode, - "v": 6, + "v": 7, "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), "dtype": str((q.dtype, ssm_backend.ssm_state.dtype)), - "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), "history": ssm_backend.capacity, "gate": int(ssm_backend.kda), "varlen": cu_seqlens is not None, @@ -56,38 +54,12 @@ def run_key(ssm_backend, q, cu_seqlens) -> tuple[int, int]: return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens -def scratch_size_bytes(ssm_backend, batch, sequences, inputs) -> int: - """估算一次 SSM autotune scratch 所需的字节数。""" - # State/history 的首轴是 layer，cursor 在各 layer 之间共享。 - layers, slots = ssm_backend.ssm_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 ) -> tuple[tuple[object, ...], dict[str, object]]: - """创建不影响服务态的独立输入，供 autotune benchmark 使用。""" + """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" 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) - # 为 benchmark 中间量、CUDA Graph 和 allocator 碎片保留一半显存。 - 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.ssm_state.shape from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig @@ -113,8 +85,7 @@ def rebuild_inputs( # benchmark 回调不能再次触发配置选择。 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) + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) ] reqs = torch.cat( ( @@ -156,7 +127,7 @@ def replay_layer(history, *args, **kwargs): @autotune( - kernel_name="ssm_update_cycle:v1", + kernel_name="ssm_update_cycle:v2", configs_gen_func=get_configs, static_key_func=static_key, run_key_func=run_key, @@ -177,7 +148,7 @@ def select_config( # 把前缀写满但不越出 capacity。只重复空 history 会测不到已有前缀的读取。 tokens = q.shape[1 if cu_seqlens is not None else 0] rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) - for step in range(rounds): + for _ in range(rounds): ssm_backend.prepare_decode(reqs, cu_seqlens) replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) ssm_backend.accept_updates(reqs, accepted) diff --git a/lightllm/common/triton_utils/autotuner.py b/lightllm/common/triton_utils/autotuner.py index c664bbe2b0..cd17bee5ff 100644 --- a/lightllm/common/triton_utils/autotuner.py +++ b/lightllm/common/triton_utils/autotuner.py @@ -164,7 +164,6 @@ def __init__( benchmark_repeats: int = 3, benchmark_retries: int = 3, ): - self.configs_gen_func = configs_gen_func self.kernel_name = kernel_name self.kernel_type = AutotuneKernelType(kernel_type) @@ -181,7 +180,6 @@ 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)} @@ -245,8 +243,6 @@ 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()) @@ -393,22 +389,7 @@ def _autotune(self, args, kwargs, static_key, run_key, rank_id, world_size): # 仅为本次调优重建输入，构造开销不计入计时，最终执行仍使用调用方的原始输入。 if self.rebuild_input_func is not None: - 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 + args, kwargs = self.rebuild_input_func(*args, **kwargs) best_config = None best_time = float("inf") 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 index 9312f663c8..726eec31c4 100644 --- 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 @@ -14,24 +14,27 @@ def ReplaySSMExecutor(*args, **kwargs): return _history_owner(*args, **kwargs) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @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 = ReplaySSMExecutor(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 = ReplaySSMExecutor(state[:1], 16, width, num_key_heads=16) - assert required == tuning.scratch_size_bytes(single_layer, 256, 256, inputs) +def test_rebuild_uses_one_layer_scratch(width): + state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + tokens = 3 * width + q = torch.ones((1, tokens, 2, 32) if width > 1 else (tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None + _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) + scratch = kwargs["ssm_backend"] + assert scratch is not ssm_backend + assert scratch.ssm_state.shape[0] == 1 + assert scratch.capacity == ssm_backend.capacity + assert torch.count_nonzero(ssm_backend.ssm_state) == 0 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("width", [1, 5]) -def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tmp_path, width): +def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType @@ -52,23 +55,24 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm 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) + + def oom(*args, **kwargs): + raise torch.cuda.OutOfMemoryError("scratch") + tuner = tuning.select_config + monkeypatch.setattr(tuner, "rebuild_input_func", oom) monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) monkeypatch.setattr(tuner, "cached_configs", {}) monkeypatch.setattr(tuner, "fast_match_configs", {}) - monkeypatch.setattr(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") + tuner, "_bench", lambda *args, **kwargs: pytest.fail("Failed rebuild must not reach 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, "num_stages": 3} + with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + tuning.resolve_run_config(**inputs) assert not list(tmp_path.iterdir()) + config = {"BV": 32, "num_warps": 1, "num_stages": 3} for name, expected in before.items(): torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py index 5eb12f9e57..547246cfe5 100644 --- a/unit_tests/common/test_autotuner.py +++ b/unit_tests/common/test_autotuner.py @@ -484,61 +484,3 @@ 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) From 8f3681f87089328447acb2d9c6596f413166c2b0 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 12:18:33 +0000 Subject: [PATCH 22/54] refactor: honor replay precompute for single-token decode Explicit compute_mode=precompute now projects a width-1 verify step, while the inline default stays. Single-token acceptance and materialize share the existing cursor and snapshot paths. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 43 +++-------- .../bounded_att/kda/replay_kda_state.py | 43 +++-------- .../basemodel/attention/linear/replay_gdn.py | 1 - .../attention/linear/replay_gdn_decode.py | 3 +- .../basemodel/attention/linear/replay_kda.py | 2 +- .../attention/linear/replay_kda_decode.py | 3 +- .../linear_att/replayssm_materialize.py | 31 +++----- .../triton_kernel/linear_att/ssm_autotune.py | 15 ++-- lightllm/server/api_cli.py | 3 +- .../common/test_linear_state_contract.py | 71 +++++++++++++++++++ 10 files changed, 116 insertions(+), 99 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 6bb90bf064..877e115fec 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -414,7 +414,7 @@ def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None if not self._is_last_layer(logical_layer_id): return - self._note_new_token(reqs) + self.accept_single_token(reqs) self.maybe_compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: @@ -432,8 +432,10 @@ def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = N del cu_seqlens return self.cursors.index_select(0, reqs.long()) - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor，保留每个请求本轮已接受的前缀。""" + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, @@ -441,14 +443,17 @@ def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tenso accepted_state_indices, reqs.numel(), self.hold_req_id, - self.verify_width, + self.verify_width if width is None else width, 256, ) def accept_single_token(self, reqs: torch.Tensor) -> None: """把一个已接受 token 记入前缀长度。""" - self._note_new_token(reqs) + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -491,7 +496,6 @@ def export_accepted_state( self, req_indices: list[int], targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, flush: bool = False, ) -> None: """将多个 accepted 快照直接写入调用方目标页。""" @@ -529,16 +533,7 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.export_accepted_state(req_indices, targets, flush=flush) return for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), - dtype=self.ssm_state.dtype, - device=self.ssm_state.device, - ) - self._materialize_accepted( - torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), - snap, - snapshot=True, - ) + snap = self.accepted_ssm(req) dst.copy_(snap) if flush: self.ssm_state[:, req].copy_(snap) @@ -574,22 +569,6 @@ def _is_last_layer(self, logical_layer_id: int) -> bool: return False return logical_layer_id == self.config.logical_layer_ids[-1] - def _note_new_token(self, reqs: torch.Tensor) -> None: - """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" - - if reqs.numel() == 0: - return - # 固定加 1，并在 kernel 内跳过 hold。避免按请求号做布尔索引，CUDA graph 不能捕获这种同步。 - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - reqs, - reqs, - reqs.numel(), - self.hold_req_id, - 1, - 256, - ) - def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 73b6253452..87fbd221d7 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -413,7 +413,7 @@ def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None if not self._is_last_layer(logical_layer_id): return - self._note_new_token(reqs) + self.accept_single_token(reqs) self.maybe_compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: @@ -431,8 +431,10 @@ def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = N del cu_seqlens return self.cursors.index_select(0, reqs.long()) - def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None) -> None: - """更新 cursor，保留每个请求本轮已接受的前缀。""" + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, @@ -440,14 +442,17 @@ def accept_updates(self, reqs: torch.Tensor, accepted_state_indices: torch.Tenso accepted_state_indices, reqs.numel(), self.hold_req_id, - self.verify_width, + self.verify_width if width is None else width, 256, ) def accept_single_token(self, reqs: torch.Tensor) -> None: """把一个已接受 token 记入前缀长度。""" - self._note_new_token(reqs) + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -490,7 +495,6 @@ def export_accepted_state( self, req_indices: list[int], targets: list[torch.Tensor], - accepted_indices: torch.Tensor | None = None, flush: bool = False, ) -> None: """将多个 accepted 快照直接写入调用方目标页。""" @@ -528,16 +532,7 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.export_accepted_state(req_indices, targets, flush=flush) return for req, dst in zip(req_indices, targets): - snap = torch.empty( - (self.ssm_state.shape[0], *self.ssm_state.shape[2:]), - dtype=self.ssm_state.dtype, - device=self.ssm_state.device, - ) - self._materialize_accepted( - torch.tensor([req], dtype=torch.int32, device=self.ssm_state.device), - snap, - snapshot=True, - ) + snap = self.accepted_ssm(req) dst.copy_(snap) if flush: self.ssm_state[:, req].copy_(snap) @@ -573,22 +568,6 @@ def _is_last_layer(self, logical_layer_id: int) -> bool: return False return logical_layer_id == self.config.logical_layer_ids[-1] - def _note_new_token(self, reqs: torch.Tensor) -> None: - """非 MTP 最后一层在写入之后，把这一个 token 记进前缀长度。""" - - if reqs.numel() == 0: - return - # 固定加 1，并在 kernel 内跳过 hold。避免按请求号做布尔索引，CUDA graph 不能捕获这种同步。 - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - reqs, - reqs, - reqs.numel(), - self.hold_req_id, - 1, - 256, - ) - def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index eb6b647e1f..db7424837b 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -180,7 +180,6 @@ def decode_att( core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) else: core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) - if self.b1_mtp_cu_q_seq_len is None: state.finish_decode_token(layer_num, self.b_conv_buffer_idx) return core, z diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index d17d8e06e7..b6348c58e6 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -36,7 +36,8 @@ def replay( axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] - precompute_projections = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index ece6f309b5..cfa3008b17 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -189,7 +189,7 @@ def decode_att( self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) - if self.b1_mtp_cu_q_seq_len is None: + if backend.mtp_step == 0: state.finish_decode_token(layer_num, self.b_conv_buffer_idx) return output diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index c4208eae31..57f085c4bf 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -36,7 +36,8 @@ def replay( axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] - precompute_projections = history.verify_width > 1 and bv >= 16 and history.compute_mode == "precompute" + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py index acf50cb5c7..a48d68b0e7 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py @@ -62,32 +62,21 @@ def _materialize( g = tl.load(Gates + slot * L + j).to(tl.float32) beta = tl.load(Betas + slot * L + j).to(tl.float32) state = _replay_update(state, raw_k, raw_v, g, beta, KDA) - if FLUSH_EXPORT: + mask = (kk[:, None] < K) & (vv[None, :] < V) + if OUTPUT_STRIDES is not None: target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) - target_out = ( + out = ( target + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + head * OUTPUT_STRIDES[1] + kk[:, None] * OUTPUT_STRIDES[2] + vv[None, :] * OUTPUT_STRIDES[3] ) - tl.store(target_out, state, (kk[:, None] < K) & (vv[None, :] < V)) - tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) - if iv == 0 and lh == 0: - tl.store(Cursors + req, 0) else: - if OUTPUT_STRIDES is not None: - target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) - out = ( - target - + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] - + head * OUTPUT_STRIDES[1] - + kk[:, None] * OUTPUT_STRIDES[2] - + vv[None, :] * OUTPUT_STRIDES[3] - ) - else: - 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)) - if WHEN_FULL and iv == 0 and lh == 0: - tl.store(Cursors + req, 0) + 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, mask) + if FLUSH_EXPORT: + tl.store(sp, state, mask) + if (FLUSH_EXPORT or WHEN_FULL) and iv == 0 and lh == 0: + tl.store(Cursors + req, 0) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 44d8d7178d..e4146aeaf1 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -24,7 +24,7 @@ def get_configs() -> list[dict[str, int]]: ] -def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens) -> dict[str, object]: +def static_key(ssm_backend, mode, q, cu_seqlens) -> dict[str, object]: """构造区分状态布局、数据类型和投影模式的静态调优键。""" _, _, hv, kd, vd = ssm_backend.ssm_state.shape # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 @@ -41,14 +41,12 @@ def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens) -> dict[str, object key["lower"] = ssm_backend.gate_lower_bound key["history_dtype"] = str(ssm_backend.keys.dtype) 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" + ("kda_mma" if ssm_backend.kda else "gdn_mma") if ssm_backend.compute_mode == "precompute" else "inline" ) return key -def run_key(ssm_backend, q, cu_seqlens) -> tuple[int, int]: +def run_key(q, cu_seqlens) -> tuple[int, int]: """返回调优使用的 token 数和请求数键。""" 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 @@ -58,7 +56,7 @@ def rebuild_inputs( ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None ) -> tuple[tuple[object, ...], dict[str, object]]: """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" - tokens, sequences = run_key(ssm_backend, q, cu_seqlens) + tokens, sequences = run_key(q, cu_seqlens) batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens _, _, hv, kd, vd = ssm_backend.ssm_state.shape from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig @@ -145,11 +143,10 @@ def select_config( ssm_backend.run_config = run_config reqs, accepted = workload ssm_backend.cursors.zero_() - # 把前缀写满但不越出 capacity。只重复空 history 会测不到已有前缀的读取。 + # 把前缀写满但不越出 capacity。replay 自己读 cursor，正式 decode 也不经过 prepare_decode。 tokens = q.shape[1 if cu_seqlens is not None else 0] rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) for _ in range(rounds): - ssm_backend.prepare_decode(reqs, cu_seqlens) replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) ssm_backend.accept_updates(reqs, accepted) return run_config @@ -160,7 +157,7 @@ def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: return ssm_backend.run_config default = _replay_config(ssm_backend.run_config) - key = (tuple(static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens).items()), run_key(ssm_backend, q, cu_seqlens)) + key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(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): diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 4965e91824..16a8735a7c 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -997,7 +997,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default={}, help='Replay-only JSON object: {"cache_len":16,"compute_mode":"precompute"}. ' "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute. " - "MTP defaults to precompute, single-token decode defaults to inline.", + "mtp_step=0 decode uses the same replay history: explicit precompute projects that one token, " + "and the default is inline. MTP defaults to precompute.", ) parser.add_argument( "--linear_att_ssm_data_type", diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index feff2d1f08..0b825cba96 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -87,6 +87,77 @@ def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") history = getattr(owner, "runtime", owner) assert history.compute_mode == expected + assert history.verify_width == mtp_step + 1 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_mtp0_decode_keeps_explicit_precompute(linear_config, cls): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + assert owner.verify_width == 1 + assert owner.compute_mode == "precompute" + assert owner.capacity == 8 + + +class _CaptureKernel: + def __init__(self): + self.args = None + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.args = args + + return launch + + +@pytest.mark.parametrize( + "module_name,cls,bv,precompute", + [ + ("replay_gdn_decode", ReplayGDNState, 32, True), + ("replay_gdn_decode", ReplayGDNState, 8, False), + ("replay_kda_decode", ReplayKDAState, 32, True), + ("replay_kda_decode", ReplayKDAState, 8, False), + ], +) +def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): + import importlib + + decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") + captured = _CaptureKernel() + monkeypatch.setattr(decode, "_replay", captured) + owner = cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + heads = owner.num_key_heads + q = torch.zeros(1, 1, heads, linear_config.key_dim) + v = torch.zeros(1, 1, owner.ssm_state.shape[-3], linear_config.value_dim) + gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] + decode.replay( + owner, + 0, + q, + q, + v, + torch.zeros(1, gate), + torch.zeros(1, v.shape[-3]), + torch.zeros(v.shape[-3]), + torch.zeros(gate), + torch.zeros(1, dtype=torch.int32), + None, + run_config={"BV": bv, "num_warps": 1, "num_stages": 1}, + ) + assert captured.args[33] is precompute + assert captured.args[35] is (cls is ReplayKDAState) @pytest.mark.parametrize("cls", [GDNState, KDAState]) From b3995226248c44d5b35273d1cde8ec7fd3f178ed Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 13:03:37 +0000 Subject: [PATCH 23/54] style: restore the blank line in Autotuner initialization The replay branch should not change autotuner.py for whitespace alone. Co-authored-by: Cursor --- lightllm/common/triton_utils/autotuner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightllm/common/triton_utils/autotuner.py b/lightllm/common/triton_utils/autotuner.py index cd17bee5ff..b7a94d657e 100644 --- a/lightllm/common/triton_utils/autotuner.py +++ b/lightllm/common/triton_utils/autotuner.py @@ -164,6 +164,7 @@ def __init__( benchmark_repeats: int = 3, benchmark_retries: int = 3, ): + self.configs_gen_func = configs_gen_func self.kernel_name = kernel_name self.kernel_type = AutotuneKernelType(kernel_type) From b1a78bb578fbd8f10b8d13eb83bc80e215788e34 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 13:07:36 +0000 Subject: [PATCH 24/54] refactor: inline one-use replay helpers Settings, the last-layer check, index upload, and prefill kernel imports only wrapped a single call site. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 38 +++++++------------ .../bounded_att/kda/replay_kda_state.py | 38 +++++++------------ .../basemodel/attention/linear/replay_gdn.py | 22 +++-------- 3 files changed, 32 insertions(+), 66 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 877e115fec..86577e4a02 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -6,7 +6,6 @@ from __future__ import annotations -from collections.abc import Sequence from functools import lru_cache from math import prod from typing import NamedTuple @@ -85,7 +84,10 @@ def __init__( 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 - capacity, compute_mode, force_compress = _replay_settings(ssm_replay_config, mtp_step) + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) _bind_history( self, self.ssm_state, @@ -406,13 +408,15 @@ def read_from_pd_page( def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层把已接受前缀整段压进 SSM。""" - if self._is_last_layer(logical_layer_id): + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: self.compress_prefix(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" - if not self._is_last_layer(logical_layer_id): + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: return self.accept_single_token(reqs) self.maybe_compress_prefix(reqs) @@ -506,9 +510,12 @@ def export_accepted_state( strides = targets[0].stride() assert len(req_indices) == len(targets) assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device self._materialize_accepted( - _upload_indices(req_indices, self.ssm_state.device, torch.int32), - _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), snapshot=not flush, flush_export=flush, output_strides=strides, @@ -564,11 +571,6 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== - def _is_last_layer(self, logical_layer_id: int) -> bool: - if self.config is None or not self.config.logical_layer_ids: - return False - return logical_layer_id == self.config.logical_layer_ids[-1] - def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 @@ -735,20 +737,6 @@ def _materialize_accepted( ) -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str, bool]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - force_compress = bool(config.get("force_compress", False)) - return capacity, compute_mode, force_compress - - -def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - def _history_owner( ssm: torch.Tensor, capacity: int, diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 87fbd221d7..a7077a98f3 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -6,7 +6,6 @@ from __future__ import annotations -from collections.abc import Sequence from functools import lru_cache from math import prod from typing import NamedTuple @@ -85,7 +84,10 @@ def __init__( 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 - capacity, compute_mode, force_compress = _replay_settings(ssm_replay_config, mtp_step) + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) _bind_history( self, self.ssm_state, @@ -405,13 +407,15 @@ def read_from_pd_page( def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层把已接受前缀整段压进 SSM。""" - if self._is_last_layer(logical_layer_id): + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: self.compress_prefix(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" - if not self._is_last_layer(logical_layer_id): + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: return self.accept_single_token(reqs) self.maybe_compress_prefix(reqs) @@ -505,9 +509,12 @@ def export_accepted_state( strides = targets[0].stride() assert len(req_indices) == len(targets) assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device self._materialize_accepted( - _upload_indices(req_indices, self.ssm_state.device, torch.int32), - _upload_indices([t.data_ptr() for t in targets], self.ssm_state.device), + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), snapshot=not flush, flush_export=flush, output_strides=strides, @@ -563,11 +570,6 @@ def run_config(self, config: dict | None) -> None: # 七、私有辅助实现 # ======================================================================== - def _is_last_layer(self, logical_layer_id: int) -> bool: - if self.config is None or not self.config.logical_layer_ids: - return False - return logical_layer_id == self.config.logical_layer_ids[-1] - def _transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 只传有效历史窗口，不包含 MTP 候选区和游标。 @@ -734,20 +736,6 @@ def _materialize_accepted( ) -def _replay_settings(config: dict | None, mtp_step: int) -> tuple[int, str, bool]: - config = dict(config or {}) - capacity = int(config.get("cache_len", 16)) - compute_mode = config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - force_compress = bool(config.get("force_compress", False)) - return capacity, compute_mode, force_compress - - -def _upload_indices(values: Sequence[int], device, dtype=torch.int64) -> torch.Tensor: - """把 Python 索引上传到目标 device。""" - - return torch.tensor(list(values), dtype=dtype, pin_memory=True).to(device, non_blocking=True) - - def _history_owner( ssm: torch.Tensor, capacity: int, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index db7424837b..a6552abeb0 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -62,7 +62,9 @@ class ReplayGDNTritonAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - _configure_gdn(self, model, _triton_prefill_kernel()) + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) def create_att_prefill_state(self, infer_state: "InferStateInfo"): return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -76,7 +78,9 @@ class ReplayGDNFlashQlaAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - _configure_gdn(self, model, _flashqla_prefill_kernel()) + from flash_qla import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) def create_att_prefill_state(self, infer_state: "InferStateInfo"): return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @@ -277,20 +281,6 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step return output -def _triton_prefill_kernel(): - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - chunk_gated_delta_rule, - ) - - return chunk_gated_delta_rule - - -def _flashqla_prefill_kernel(): - from flash_qla import chunk_gated_delta_rule - - return chunk_gated_delta_rule - - def _configure_gdn(backend, model, prefill_kernel) -> None: """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" From a8cb33b108b68c1905722a4cf88f84037145a408 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 13:15:32 +0000 Subject: [PATCH 25/54] refactor: put replay entry points ahead of their helpers Readers now meet the decode kernel, autotune resolution, and accepted-tensor type before the helpers they call. Decode init still requires exactly one replay state. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 10 +- .../bounded_att/kda/replay_kda_state.py | 10 +- .../basemodel/attention/linear/replay_gdn.py | 41 +- .../basemodel/attention/linear/replay_kda.py | 23 +- .../triton_kernel/linear_att/replayssm.py | 8 +- .../linear_att/replayssm_compute.py | 398 +++++++++--------- .../triton_kernel/linear_att/ssm_autotune.py | 34 +- 7 files changed, 257 insertions(+), 267 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 86577e4a02..7c8db6a5d4 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -31,6 +31,11 @@ logger = init_logger(__name__) +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -819,8 +824,3 @@ def _bind_history( owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) - - -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index a7077a98f3..7f9e68bd4b 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -31,6 +31,11 @@ logger = init_logger(__name__) +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -818,8 +823,3 @@ def _bind_history( owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) - - -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index a6552abeb0..d08a714d27 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -158,9 +158,16 @@ class ReplayGDNDecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def init_state(self): - state = _single_replay_state(self) + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayGDNState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}") draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - self._init_indices(draft_step, state.mtp_state_index) + self._init_indices(draft_step, matches[0].mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx def decode_att( @@ -281,6 +288,15 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step return output +def _replay_state(att_state, layer_num: int) -> ReplayGDNState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=ReplayGDNState, + ) + + def _configure_gdn(backend, model, prefill_kernel) -> None: """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" @@ -334,24 +350,3 @@ def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) return query, key, value - - -def _replay_state(att_state, layer_num: int) -> ReplayGDNState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=ReplayGDNState, - ) - - -def _single_replay_state(att_state) -> ReplayGDNState: - role = att_state.infer_state.att_state_role - matches = [ - state - for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayGDNState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayGDNState for {role}") - return matches[0] diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index cfa3008b17..3090faa9a8 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -135,8 +135,15 @@ class ReplayKDADecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def init_state(self): - state = _single_replay_state(self) - self._init_indices(self.backend.mtp_step, state.mtp_state_index) + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayKDAState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}") + self._init_indices(self.backend.mtp_step, matches[0].mtp_state_index) self.b_ssm_buffer_idx = self.b_conv_buffer_idx def decode_att( @@ -232,15 +239,3 @@ def _replay_state(att_state, layer_num: int) -> ReplayKDAState: logical_layer_id=layer_num, state_class=ReplayKDAState, ) - - -def _single_replay_state(att_state) -> ReplayKDAState: - role = att_state.infer_state.att_state_role - matches = [ - state - for state in att_state.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayKDAState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayKDAState for {role}") - return matches[0] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index a8751ea970..74e7b973e9 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -5,22 +5,22 @@ existing module-level names used by the executor and tests. """ -from .replayssm_common import _replay_update from .replayssm_compute import ( _project_replay_queries, _replay, _replay_gdn_window, _replay_sequential, ) +from .replayssm_common import _replay_update from .replayssm_fold import _accept_updates from .replayssm_materialize import _materialize __all__ = [ - "_accept_updates", - "_materialize", - "_project_replay_queries", "_replay", + "_project_replay_queries", "_replay_gdn_window", "_replay_sequential", "_replay_update", + "_accept_updates", + "_materialize", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py index c40c5bf44f..18ea024e8a 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py @@ -5,6 +5,205 @@ import triton.language as tl +@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, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, +): + """执行一个请求窗口的 replay decode，并写入 history。 + + 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 + cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， + 不应因为 eager batch 改变而形成新的 kernel 变体。 + """ + 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) + n = tl.load(Cursors + req) + 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 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" + qk, projections, tt, tokens, valid, h = _project_replay_queries( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ, + SK, + SA, + H, + HV, + K, + L, + BW, + KDA, + LOWER, + BQK, + end, + hv, + kk, + ll, + n, + projection_dtype, + projection_precision, + slot, + start, + state, + ) + else: + projections, tt = None, None + + if PRECOMPUTE_PROJECTIONS and not KDA: + _replay_gdn_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype, + projection_precision, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, + ) + else: + _replay_sequential( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ, + SK, + SV, + SA, + SB, + H, + HV, + KH, + K, + V, + L, + PRECOMPUTE_PROJECTIONS, + BW, + KDA, + LOWER, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, + ) + + @triton.jit def _project_replay_queries( Q, @@ -351,202 +550,3 @@ def _replay_sequential( gate_prefix = tl.where(ll >= n, total_g, gate_prefix) history |= ll == n n += 1 - - -@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, - KDA: tl.constexpr, - LOWER: tl.constexpr, - BQK: tl.constexpr, - BR: tl.constexpr, -): - """执行一个请求窗口的 replay decode，并写入 history。 - - 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 - cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， - 不应因为 eager batch 改变而形成新的 kernel 变体。 - """ - 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) - n = tl.load(Cursors + req) - 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 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" - qk, projections, tt, tokens, valid, h = _project_replay_queries( - Q, - Kp, - A, - Alog, - Bias, - Gates, - SQ, - SK, - SA, - H, - HV, - K, - L, - BW, - KDA, - LOWER, - BQK, - end, - hv, - kk, - ll, - n, - projection_dtype, - projection_precision, - slot, - start, - state, - ) - else: - projections, tt = None, None - - if PRECOMPUTE_PROJECTIONS and not KDA: - _replay_gdn_window( - Kp, - Vp, - A, - B, - Alog, - Bias, - Keys, - Deltas, - Gates, - RawKeys, - RawValues, - Betas, - Out, - SK, - SV, - SA, - SB, - HV, - KH, - K, - V, - L, - BV, - BW, - BR, - end, - h, - hv, - iv, - key_slot, - kk, - n, - projection_dtype, - projection_precision, - projections, - qk, - slot, - start, - tokens, - tt, - valid, - vv, - ) - else: - _replay_sequential( - Q, - Kp, - Vp, - A, - B, - Alog, - Bias, - Keys, - Deltas, - Gates, - RawKeys, - RawValues, - Betas, - Out, - SQ, - SK, - SV, - SA, - SB, - H, - HV, - KH, - K, - V, - L, - PRECOMPUTE_PROJECTIONS, - BW, - KDA, - LOWER, - end, - hv, - iv, - key_slot, - kk, - ll, - n, - projections, - slot, - start, - state, - tt, - vv, - ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index e4146aeaf1..3a25fd5f9f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -15,6 +15,23 @@ _DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: + """返回当前 batch 的固定或已调优 SSM kernel 配置。""" + if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: + return ssm_backend.run_config + default = _replay_config(ssm_backend.run_config) + key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(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 = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) + # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 + ssm_backend.batch_configs[key] = config + logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") + return config + + def get_configs() -> list[dict[str, int]]: """返回 ReplaySSM update cycle 的候选 Triton 配置。""" return [ @@ -152,23 +169,6 @@ def select_config( return run_config -def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: - """返回当前 batch 的固定或已调优 SSM kernel 配置。""" - if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: - return ssm_backend.run_config - default = _replay_config(ssm_backend.run_config) - key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(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 = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) - # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 - ssm_backend.batch_configs[key] = config - logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") - return config - - def _replay_config(config: dict | None) -> dict[str, int]: if not config: return dict(_DEFAULT_REPLAY_CONFIG) From 833cf515fbed782a791171b1ba64cd5afcc69f14 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 13:22:49 +0000 Subject: [PATCH 26/54] refactor: restore the original MTP state-params layout Replay reuses the existing dynamic verify helper, so this file does not need its kernel moved. Co-authored-by: Cursor --- .../linear_att/mtp_state_params.py | 88 +++++++++---------- 1 file changed, 44 insertions(+), 44 deletions(-) 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 266e2d2b49..b81bf9ec16 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 @@ -3,6 +3,50 @@ import triton.language as tl +@triton.jit +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, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.arange(0, BLOCK_SIZE) + token_mask = offsets < batch_size + cu_mask = offsets <= batch_size + + 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) + + 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. + 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.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, @@ -65,47 +109,3 @@ def build_dynamic_mtp_linear_att_state_params( num_stages=1, ) return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens - - -@triton.jit -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, - BLOCK_SIZE: tl.constexpr, -): - offsets = tl.arange(0, BLOCK_SIZE) - token_mask = offsets < batch_size - cu_mask = offsets <= batch_size - - 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) - - 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. - 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.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) From eb91e421a4f667ad0971e8467f7bd258bea5fc0c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 9 Nov 2026 13:32:49 +0000 Subject: [PATCH 27/54] refactor: name replay kernels by decode, cursor, and fold The old fold file only advanced the cursor, while the state update lived under materialize. File and function names now match those steps. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 8 +- .../bounded_att/kda/replay_kda_state.py | 8 +- .../attention/linear/replay_gdn_decode.py | 4 +- .../attention/linear/replay_kda_decode.py | 4 +- .../triton_kernel/linear_att/replayssm.py | 37 ++++---- .../linear_att/replayssm_cursor.py | 24 +++++ ...playssm_compute.py => replayssm_decode.py} | 21 +++-- ...replayssm_common.py => replayssm_delta.py} | 6 +- .../linear_att/replayssm_fold.py | 90 +++++++++++++++---- .../linear_att/replayssm_materialize.py | 82 ----------------- .../linear_att/test_replayssm_execution.py | 6 +- .../common/test_linear_state_contract.py | 2 +- 12 files changed, 149 insertions(+), 143 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py rename lightllm/common/basemodel/triton_kernel/linear_att/{replayssm_compute.py => replayssm_decode.py} (97%) rename lightllm/common/basemodel/triton_kernel/linear_att/{replayssm_common.py => replayssm_delta.py} (71%) delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 7c8db6a5d4..3b9113db18 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -14,8 +14,8 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( - _accept_updates, - _materialize, + _advance_cursor, + _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -446,7 +446,7 @@ def accept_updates( ) -> None: """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, reqs, accepted_state_indices, @@ -713,7 +713,7 @@ def _materialize_accepted( layers, slots, hv, k, v = self.ssm_state.shape config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} bv = config["BV"] - _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, self.raw_keys, diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 7f9e68bd4b..bc3b2b9c1c 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -14,8 +14,8 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( - _accept_updates, - _materialize, + _advance_cursor, + _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -445,7 +445,7 @@ def accept_updates( ) -> None: """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, reqs, accepted_state_indices, @@ -712,7 +712,7 @@ def _materialize_accepted( layers, slots, hv, k, v = self.ssm_state.shape config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} bv = config["BV"] - _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, self.raw_keys, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index b6348c58e6..71fcf8e1aa 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -8,7 +8,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config @@ -40,7 +40,7 @@ def replay( precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) - _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, v, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 57f085c4bf..55514aeb57 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -8,7 +8,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config @@ -40,7 +40,7 @@ def replay( precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) - _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, v, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 74e7b973e9..dcda881c0c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -1,26 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 -"""Public ReplaySSM kernel exports. +"""ReplaySSM kernel exports, in lifecycle order. -Kernel implementations are split by lifecycle stage; these imports preserve the -existing module-level names used by the executor and tests. +decode writes history for the current window. +cursor records how many of those tokens were accepted. +fold applies that accepted prefix onto the canonical SSM. """ -from .replayssm_compute import ( - _project_replay_queries, - _replay, - _replay_gdn_window, - _replay_sequential, +from .replayssm_decode import ( + _decode_replay, + _gdn_precomputed_window, + _project_qk, + _replay_each_token, ) -from .replayssm_common import _replay_update -from .replayssm_fold import _accept_updates -from .replayssm_materialize import _materialize +from .replayssm_delta import _delta_rule_step +from .replayssm_cursor import _advance_cursor +from .replayssm_fold import _fold_prefix __all__ = [ - "_replay", - "_project_replay_queries", - "_replay_gdn_window", - "_replay_sequential", - "_replay_update", - "_accept_updates", - "_materialize", + "_decode_replay", + "_project_qk", + "_gdn_precomputed_window", + "_replay_each_token", + "_delta_rule_step", + "_advance_cursor", + "_fold_prefix", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py new file mode 100644 index 0000000000..5354c71260 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Advance the accepted-prefix cursor. This kernel does not fold state.""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _advance_cursor( + cursors, + reqs, + accepted_state_indices, + N, + HOLD: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 + 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) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py rename to lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 18ea024e8a..25ef5d0fdf 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compute.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -"""ReplaySSM projection and recurrent decode kernels.""" +"""ReplaySSM decode kernels. + +`_decode_replay` is the entry. It projects the verify window when asked, then +either solves a GDN window or walks tokens one by one. Neither path writes the +canonical SSM. +""" import triton import triton.language as tl @triton.jit -def _replay( +def _decode_replay( Q, Kp, Vp, @@ -80,7 +85,7 @@ def _replay( 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" - qk, projections, tt, tokens, valid, h = _project_replay_queries( + qk, projections, tt, tokens, valid, h = _project_qk( Q, Kp, A, @@ -113,7 +118,7 @@ def _replay( projections, tt = None, None if PRECOMPUTE_PROJECTIONS and not KDA: - _replay_gdn_window( + _gdn_precomputed_window( Kp, Vp, A, @@ -158,7 +163,7 @@ def _replay( vv, ) else: - _replay_sequential( + _replay_each_token( Q, Kp, Vp, @@ -205,7 +210,7 @@ def _replay( @triton.jit -def _project_replay_queries( +def _project_qk( Q, Kp, A, @@ -287,7 +292,7 @@ def _project_replay_queries( @triton.jit -def _replay_gdn_window( +def _gdn_precomputed_window( Kp, Vp, A, @@ -419,7 +424,7 @@ def _replay_gdn_window( @triton.jit -def _replay_sequential( +def _replay_each_token( Q, Kp, Vp, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py similarity index 71% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py rename to lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py index 84c1af9abe..b3a6b2aec0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_common.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -"""Shared ReplaySSM recurrence primitives.""" +"""One gated-delta step shared by prefix folding.""" import triton import triton.language as tl @triton.jit -def _replay_update(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): - # Replay folding and materialization use the same update order. +def _delta_rule_step(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + # Folding applies this same update once per accepted history token. raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) decay = tl.exp(g) if KDA: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 136ead33f8..f908916bc7 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -1,24 +1,82 @@ # SPDX-License-Identifier: Apache-2.0 -"""ReplaySSM accepted-prefix length updates.""" +"""Fold an accepted prefix into the canonical SSM.""" import triton import triton.language as tl +from .replayssm_delta import _delta_rule_step -@triton.jit(do_not_specialize=["N"]) -def _accept_updates( - cursors, - reqs, - accepted_state_indices, - N, + +@triton.jit +def _fold_prefix( + 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, - WIDTH: tl.constexpr, - BLOCK: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SNAPSHOT: tl.constexpr, + KDA: tl.constexpr, + FLUSH_EXPORT: tl.constexpr = False, + OUTPUT_STRIDES: tl.constexpr = None, + WHEN_FULL: tl.constexpr = False, + NEXT: tl.constexpr = 1, ): - # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 - 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) + # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT + # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， + # 并在 kernel 内把这些请求的 cursor 清零，避免主机分支进不了 CUDA Graph。 + 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 + n = tl.load(Cursors + req) + if WHEN_FULL: + if n == 0 or n + NEXT <= L: + return + elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + 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 * L + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, vv < V, 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, kk < K, 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j).to(tl.float32) + beta = tl.load(Betas + slot * L + j).to(tl.float32) + state = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + mask = (kk[:, None] < K) & (vv[None, :] < V) + if OUTPUT_STRIDES is not None: + target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) + out = ( + target + + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + + head * OUTPUT_STRIDES[1] + + kk[:, None] * OUTPUT_STRIDES[2] + + vv[None, :] * OUTPUT_STRIDES[3] + ) + else: + 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, mask) + if FLUSH_EXPORT: + tl.store(sp, state, mask) + if (FLUSH_EXPORT or WHEN_FULL) and iv == 0 and lh == 0: + tl.store(Cursors + req, 0) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py deleted file mode 100644 index a48d68b0e7..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_materialize.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""ReplaySSM accepted-state materialization kernel.""" - -import triton -import triton.language as tl - -from .replayssm_common import _replay_update - - -@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, - FLUSH_EXPORT: tl.constexpr = False, - OUTPUT_STRIDES: tl.constexpr = None, - WHEN_FULL: tl.constexpr = False, - NEXT: tl.constexpr = 1, -): - # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT - # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， - # 并在 kernel 内把这些请求的 cursor 清零，避免主机分支进不了 CUDA Graph。 - 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 - n = tl.load(Cursors + req) - if WHEN_FULL: - if n == 0 or n + NEXT <= L: - return - elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: - 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 * L + j) * K + kk, kk < K, 0).to(tl.float32) - raw_v = tl.load(RawValues + (slot * L + j) * V + vv, vv < V, 0).to(tl.float32) - if KDA: - g = tl.load(Gates + (slot * L + j) * K + kk, kk < K, 0).to(tl.float32) - else: - g = tl.load(Gates + slot * L + j).to(tl.float32) - beta = tl.load(Betas + slot * L + j).to(tl.float32) - state = _replay_update(state, raw_k, raw_v, g, beta, KDA) - mask = (kk[:, None] < K) & (vv[None, :] < V) - if OUTPUT_STRIDES is not None: - target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) - out = ( - target - + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] - + head * OUTPUT_STRIDES[1] - + kk[:, None] * OUTPUT_STRIDES[2] - + vv[None, :] * OUTPUT_STRIDES[3] - ) - else: - 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, mask) - if FLUSH_EXPORT: - tl.store(sp, state, mask) - if (FLUSH_EXPORT or WHEN_FULL) and iv == 0 and lh == 0: - tl.store(Cursors + req, 0) 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 index 5ca0d1186b..7ab9d301f8 100644 --- 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 @@ -25,7 +25,7 @@ def ReplaySSMExecutor(*args, **kwargs): def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): from lightllm.common.basemodel.triton_kernel.linear_att import replayssm - kernels = ("_replay", "_accept_updates") + kernels = ("_decode_replay", "_advance_cursor") hashes = {name: set() for name in kernels} launches = [] for name in kernels: @@ -86,8 +86,8 @@ def cycle(backend, ids): torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) assert torch.equal(captured.cursors, eager.cursors) torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) - assert len(hashes["_replay"]) == 1 - assert len(hashes["_accept_updates"]) == 1 + assert len(hashes["_decode_replay"]) == 1 + assert len(hashes["_advance_cursor"]) == 1 assert captured.cursors[-1] == 0 assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 0b825cba96..f720e01f09 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -130,7 +130,7 @@ def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, modul decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") captured = _CaptureKernel() - monkeypatch.setattr(decode, "_replay", captured) + monkeypatch.setattr(decode, "_decode_replay", captured) owner = cls( linear_config, request_capacity=2, From 2f117c67b96aeb5fa726a5fdb8dc5f4e9d2a7fcd Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Tue, 10 Nov 2026 06:57:08 +0000 Subject: [PATCH 28/54] feat: reconstruct GDN and KDA replay outputs on tensor cores Use tf32x3 dots and an fp16 k/d cache for bf16 activations, and give KDA the same windowed reconstruction. Folding stays an fp32 pass over raw k/v, and history loads stay inside the prefix. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 13 +- .../bounded_att/kda/replay_kda_state.py | 13 +- lightllm/common/attention_state/manager.py | 12 +- .../triton_kernel/linear_att/replayssm.py | 2 + .../linear_att/replayssm_decode.py | 319 +++++++++++++++--- .../linear_att/replayssm_fold.py | 17 +- .../linear_att/test_replayssm.py | 47 +++ .../linear_att/test_replayssm_execution.py | 37 ++ 8 files changed, 389 insertions(+), 71 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 3b9113db18..b4f03b2602 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -811,15 +811,12 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - # 一段连续前缀，长度不超过 capacity。precompute 且 BF16、短 history 时 history 用激活精度，其余用 FP32。 - history_dtype = ( - activation_dtype - if activation_dtype == torch.bfloat16 and verify_width > 1 and compute_mode == "precompute" and capacity <= 8 - else torch.float32 - ) + # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) - owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=history_dtype, device=ssm.device) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) owner.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index bc3b2b9c1c..7df2df8738 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -810,15 +810,12 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - # 一段连续前缀，长度不超过 capacity。precompute 且 BF16、短 history 时 history 用激活精度，其余用 FP32。 - history_dtype = ( - activation_dtype - if activation_dtype == torch.bfloat16 and verify_width > 1 and compute_mode == "precompute" and capacity <= 8 - else torch.float32 - ) + # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) - owner.keys = torch.empty(key_shape, dtype=history_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=history_dtype, device=ssm.device) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) owner.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=ssm.device) owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6f235d3a6d..6fb18c7ea1 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -59,6 +59,12 @@ def __init__(self, request_capacity: int): self._dp_shared_operator = DPSharedOperator(self) self._pd_operator = PDOperator(self) + def _reset_operators(self) -> None: + """让浅拷贝后的 Operator 指向新的资源管理器副本。""" + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + def get_cpu_cache_operator(self) -> CpuCacheOperator: """返回当前资源管理器唯一的 CPU cache 操作对象。""" return self._cpu_cache_operator @@ -406,12 +412,6 @@ def commit_bounded_att_state( # 七、私有辅助实现 # ======================================================================== - def _reset_operators(self) -> None: - """让浅拷贝后的 Operator 指向新的资源管理器副本。""" - self._cpu_cache_operator = CpuCacheOperator(self) - self._dp_shared_operator = DPSharedOperator(self) - self._pd_operator = PDOperator(self) - # 私有实现：缓存创建完成后的成员集合，模型逐层计算直接访问具体 State 成员。 @lru_cache(maxsize=None) def _get_states( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index dcda881c0c..e54186eae4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -9,6 +9,7 @@ from .replayssm_decode import ( _decode_replay, _gdn_precomputed_window, + _kda_precomputed_window, _project_qk, _replay_each_token, ) @@ -20,6 +21,7 @@ "_decode_replay", "_project_qk", "_gdn_precomputed_window", + "_kda_precomputed_window", "_replay_each_token", "_delta_rule_step", "_advance_cursor", diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 25ef5d0fdf..5619869b64 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -2,7 +2,7 @@ """ReplaySSM decode kernels. `_decode_replay` is the entry. It projects the verify window when asked, then -either solves a GDN window or walks tokens one by one. Neither path writes the +solves a GDN or KDA window, or walks tokens one by one. Neither path writes the canonical SSM. """ @@ -10,6 +10,22 @@ import triton.language as tl +@triton.jit +def _history_coeff(rows, vec, BK: tl.constexpr): + # rows [R, BK] · vec [BK] -> [R]. tf32x3 是 tensor core 上最接近 fp32 的输出重建。 + if BK >= 16: + return tl.sum(tl.dot(rows, vec[:, None], input_precision="tf32x3"), 1) + return tl.sum(rows * vec[None, :], 1) + + +@triton.jit +def _apply_deltas(coeff, deltas, R: tl.constexpr, BV: tl.constexpr): + # coeff [R] · deltas [R, BV] -> [BV]. R 小于 16 时 tensor core 不能做这次点积。 + if R >= 16 and BV >= 16: + return tl.sum(tl.dot(coeff[None, :], deltas, input_precision="tf32x3"), 0) + return tl.sum(deltas * coeff[:, None], 0) + + @triton.jit def _decode_replay( Q, @@ -81,10 +97,10 @@ def _decode_replay( state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) 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" + # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 + # 状态不在这里更新；压缩仍用原始 k/v 做 fp32 串行折回。 + projection_dtype: tl.constexpr = tl.float32 + projection_precision: tl.constexpr = "tf32x3" qk, projections, tt, tokens, valid, h = _project_qk( Q, Kp, @@ -117,7 +133,50 @@ def _decode_replay( else: projections, tt = None, None - if PRECOMPUTE_PROJECTIONS and not KDA: + if PRECOMPUTE_PROJECTIONS and KDA: + _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + BK, + LOWER, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, + ) + elif PRECOMPUTE_PROJECTIONS: _gdn_precomputed_window( Kp, Vp, @@ -193,6 +252,8 @@ def _decode_replay( BW, KDA, LOWER, + BK, + BV, end, hv, iv, @@ -243,20 +304,25 @@ def _project_qk( tt = tl.arange(0, BQK) tokens = start + tt % BW valid = (tt < 2 * BW) & (tokens < end) + safe_tokens = tl.where(valid, tokens, start) h = hv // (HV // H) q_rows = tl.load( - Q + tokens[:, None] * SQ + h * K + kk[None, :], + Q + safe_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, :], + Kp + safe_tokens[:, None] * SK + h * K + kk[None, :], (valid & (tt >= BW))[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) - qk = q_rows + k_rows + # Padded query/key rows share the MMA tile. Keep them at zero so they + # cannot change the valid rows. + row_valid = valid[:, None] + qk = tl.where(row_valid, q_rows + k_rows, 0) qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk = tl.where(row_valid, qk, 0) if KDA: history = ll < n gates = tl.load( @@ -266,9 +332,11 @@ def _project_qk( ).to(tl.float32) total_g = tl.sum(gates, axis=0) wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) gate_inputs = tl.load( - A + (start + wi[:, None]) * SA + hv * K + kk[None, :], - (start + wi[:, None] < end) & (kk[None, :] < K), + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] @@ -281,8 +349,10 @@ def _project_qk( ) # All exponents are forward decay. No inverse prefix scaling can # overflow when a long history contains strongly negative gates. + # Padded rows stay zero: 0 * exp(garbage) is NaN and would enter the dot. + scaled = tl.where(row_valid, qk * tl.exp(query_gates), 0) projections = tl.dot( - (qk * tl.exp(query_gates)).to(projection_dtype), + scaled.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision, ) @@ -342,27 +412,32 @@ def _gdn_precomputed_window( fresh_tokens = start + jj - L old = (jj < L) & (jj < n) fresh = (jj >= L) & (jj < L + end - start) + # BR is rounded above the history length. Masked lanes must still address + # a real slot; otherwise they read uninitialized neighbor rows into the dot. + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) cached_keys = tl.load( - Keys + (key_slot * L) * K + jj[:, None] * K + kk[None, :], + Keys + (key_slot * L) * K + safe_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, :], + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], fresh[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) all_keys = cached_keys + fresh_keys values = tl.load( - Vp + fresh_tokens[:, None] * SV + hv * V + vv[None, :], + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], fresh[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) - x = tl.load(A + fresh_tokens * SA + hv, fresh, 0).to(tl.float32) + bias + x = tl.load(A + safe_fresh * SA + hv, fresh, 0).to(tl.float32) + bias fresh_gates = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) - gates = tl.load(Gates + slot * L + jj, old, 0).to(tl.float32) + gates = tl.load(Gates + slot * L + safe_jj, old, 0).to(tl.float32) gates += tl.where(fresh, fresh_gates, 0.0) prefix = tl.cumsum(gates, axis=0) query_gates = tl.sum(tl.where(jj[None, :] <= L + (tt % BW)[:, None], gates[None, :], 0), 1) @@ -373,15 +448,15 @@ def _gdn_precomputed_window( coefficients *= tl.where(causal, tl.exp(tl.minimum(query_gates[:, None] - prefix[None, :], 0.0)), 0.0) projections *= tl.exp(query_gates)[:, None] deltas = tl.load( - Deltas + (slot * L) * V + jj[:, None] * V + vv[None, :], + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], old[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) - betas = tl.sigmoid(tl.load(B + fresh_tokens * SB + hv, fresh, 0).to(tl.float32)) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) for t in range(end - start): checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) - sk = checkpoint_sk + tl.sum(deltas * weights[:, None], 0) + sk = checkpoint_sk + _apply_deltas(weights, deltas, BR, BV) value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) beta = tl.sum(tl.where(jj == L + t, betas, 0)) delta = beta * (value - sk) @@ -389,37 +464,193 @@ def _gdn_precomputed_window( outputs = projections + tl.dot( coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision ) + out_tokens = tl.where(valid, tokens, start) tl.store( - Out + (tokens[:, None] * HV + hv) * V + vv[None, :], + Out + (out_tokens[:, None] * HV + hv) * V + vv[None, :], outputs, (valid & (tt < BW))[:, None] & (vv[None, :] < V), ) records = n + jj - L in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Gates + slot * L + safe_records, fresh_gates, in_prefix) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + BK: tl.constexpr, + LOWER: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, +): + # KDA 的衰减在每个 key 通道上，不能把标量 exp 提到点积外面。 + # 每个 token 用冻结 checkpoint 的投影，加上 history 的通道衰减点积重建输出。 + # exp(G_t - G_j) 对已可见的 token 不大于 1。新 delta 仍按 token 解。 + # 状态压缩继续用原始 k/v 在 fp32 里折回。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = tl.where(old[:, None], cached_keys, 0) + fresh_keys + values = tl.load( + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + hist_g = tl.load( + Gates + (slot * L + safe_jj[:, None]) * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_hist = tl.sum(hist_g, axis=0) + hist_prefix = tl.cumsum(hist_g, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_g_rows = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv).to(tl.float32)) * gate_inputs) + fresh_g_rows = tl.where(gate_ok[:, None], fresh_g_rows, 0) + fresh_prefix = tl.cumsum(fresh_g_rows, axis=0) + fresh_index = tl.broadcast_to(tl.where(fresh, jj - L, 0)[:, None], (BR, BK)) + G_pos = tl.where(old[:, None], hist_prefix, 0.0) + G_pos = tl.where(fresh[:, None], total_hist[None, :] + tl.gather(fresh_prefix, fresh_index, 0), G_pos) + step_g = tl.where(fresh[:, None], tl.gather(fresh_g_rows, fresh_index, 0), 0.0) + G_query = total_hist[None, :] + tl.gather(fresh_prefix, tl.broadcast_to((tt % BW)[:, None], qk.shape), 0) + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(start, end): + local = t - start + q = tl.sum(tl.where((tt == local)[:, None], qk, 0), 0) + k = tl.sum(tl.where((tt == BW + local)[:, None], qk, 0), 0) + g_now = tl.sum(tl.where((tt == local)[:, None], G_query, 0), 0) + causal = old | (fresh & (jj <= L + local)) + decay = tl.exp(tl.where(causal[:, None], g_now[None, :] - G_pos, 0.0)) + scaled = tl.where(causal[:, None], all_keys * decay, 0) + hk = _history_coeff(scaled, k, BK) + proj_k = tl.sum(tl.where((tt == BW + local)[:, None], projections, 0), 0) + sk = proj_k + _apply_deltas(hk, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + local)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + local, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + local)[:, None], delta[None, :], deltas) + hq = _history_coeff(scaled, q, BK) + proj_q = tl.sum(tl.where((tt == local)[:, None], projections, 0), 0) + sq = proj_q + _apply_deltas(hq, deltas, BR, BV) + tl.store(Out + (t * HV + hv) * V + vv, sq, vv < V) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) tl.store( - Deltas + (slot * L + records[:, None]) * V + vv[None, :], + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], deltas, in_prefix[:, None] & (vv[None, :] < V), ) tl.store( - RawValues + (slot * L + records[:, None]) * V + vv[None, :], + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], values, in_prefix[:, None] & (vv[None, :] < V), ) if iv == 0: if hv % (HV // KH) == 0: tl.store( - Keys + (key_slot * L + records[:, None]) * K + kk[None, :], + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], fresh_keys, in_prefix[:, None] & (kk[None, :] < K), ) tl.store( - RawKeys + (key_slot * L + records[:, None]) * K + kk[None, :], + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], raw_keys, in_prefix[:, None] & (kk[None, :] < K), ) - tl.store(Gates + slot * L + records, fresh_gates, in_prefix) - tl.store(Betas + slot * L + records, betas, in_prefix) + tl.store( + Gates + (slot * L + safe_records[:, None]) * K + kk[None, :], + step_g, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) return @@ -454,6 +685,8 @@ def _replay_each_token( BW: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, end, hv, iv, @@ -514,41 +747,43 @@ def _replay_each_token( # 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) + hk = _history_coeff(keys, k, BK) + hq = _history_coeff(keys, q, BK) 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) + sk = checkpoint_sk + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq + _apply_deltas(hq, deltas, L, BV) else: weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) - hk = tl.sum(keys * k[None, :], axis=1) * weights - hq = tl.sum(keys * q[None, :], axis=1) * weights + hk = _history_coeff(keys, k, BK) * weights + hq = _history_coeff(keys, q, BK) * weights checkpoint_sk = tl.sum(state * k[:, None], axis=0) checkpoint_sq = tl.sum(state * q[:, None], axis=0) - sk = checkpoint_sk * tl.exp(total_g) + tl.sum(deltas * hk[:, None], axis=0) - sq = checkpoint_sq * tl.exp(total_g) + tl.sum(deltas * hq[:, None], axis=0) + sk = checkpoint_sk * tl.exp(total_g) + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq * tl.exp(total_g) + _apply_deltas(hq, deltas, L, BV) d = beta * (v - sk) tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) record = n in_prefix = record < L - tl.store(Deltas + (slot * L + record) * V + vv, d, in_prefix & (vv < V)) - tl.store(RawValues + (slot * L + record) * V + vv, v, in_prefix & (vv < V)) + # 被 mask 掉的 lane 仍会计算地址。cursor 已经顶到容量时必须改写 0 号槽，避免非法访问。 + safe_record = tl.where(in_prefix, record, 0) + tl.store(Deltas + (slot * L + safe_record) * V + vv, d, in_prefix & (vv < V)) + tl.store(RawValues + (slot * L + safe_record) * V + vv, v, in_prefix & (vv < V)) if iv == 0: # Only the group owner writes shared K. Every CTA computes this # round's K locally; readers only consume earlier accepted records. if hv % (HV // KH) == 0: - tl.store(Keys + (key_slot * L + record) * K + kk, k, in_prefix & (kk < K)) - tl.store(RawKeys + (key_slot * L + record) * K + kk, raw_k, in_prefix & (kk < K)) + tl.store(Keys + (key_slot * L + safe_record) * K + kk, k, in_prefix & (kk < K)) + tl.store(RawKeys + (key_slot * L + safe_record) * K + kk, raw_k, in_prefix & (kk < K)) if KDA: - tl.store(Gates + (slot * L + record) * K + kk, g, in_prefix & (kk < K)) + tl.store(Gates + (slot * L + safe_record) * K + kk, g, in_prefix & (kk < K)) else: - tl.store(Gates + slot * L + record, g, in_prefix) - tl.store(Betas + slot * L + record, beta, in_prefix) + tl.store(Gates + slot * L + safe_record, g, in_prefix) + tl.store(Betas + slot * L + safe_record, beta, in_prefix) keys = tl.where((ll == n)[:, None], k[None, :], keys) deltas = tl.where((ll == n)[:, None], d[None, :], deltas) if not KDA: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index f908916bc7..1a10d122dd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -53,15 +53,18 @@ def _fold_prefix( 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 * L + j) * K + kk, kk < K, 0).to(tl.float32) - raw_v = tl.load(RawValues + (slot * L + j) * V + vv, vv < V, 0).to(tl.float32) + # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 + for j in range(L): + active = j < n + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) if KDA: - g = tl.load(Gates + (slot * L + j) * K + kk, kk < K, 0).to(tl.float32) + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) else: - g = tl.load(Gates + slot * L + j).to(tl.float32) - beta = tl.load(Betas + slot * L + j).to(tl.float32) - state = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) mask = (kk[:, None] < K) & (vv[None, :] < V) if OUTPUT_STRIDES is not None: target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) 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 index 9aab192432..23270bb407 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -395,6 +395,53 @@ def test_bf16_commit_rounds_before_the_next_token(): torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) +def test_large_bf16_state_precompute_matches_inline(): + # 输出投影用 tf32x3。重建缓存的 k/d 在 bf16 激活下是 fp16，输出允许一个 fp16 单位。 + # 折回 SSM 仍从原始 k/v 在 fp32 中计算，两条路径的 checkpoint 必须一致。 + heads = 16 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros_like(alog) + + def check(capacity, width, steps): + torch.manual_seed(7) + state = (torch.randn((1, 2, heads, 128, 128), device="cuda") * 10).bfloat16() + backends = [ + ReplaySSMExecutor( + state.clone(), + capacity, + width, + torch.bfloat16, + num_key_heads=heads, + compute_mode=mode, + run_config=config, + ) + for mode in ("inline", "precompute") + ] + cu = None if width == 1 else torch.tensor([0, width], dtype=torch.int32, device="cuda") + accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") + for step in range(steps): + torch.manual_seed(100 + step) + q = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 + b = torch.randn_like(a) + outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] for backend in backends] + # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 + torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) + for backend in backends: + backend.accept_updates(reqs, accepted, width=width) + backend.maybe_compress_prefix(reqs) + for backend in backends: + backend.merge_accepted_updates(reqs) + torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) + + check(4, 1, 16) + check(16, 2, 4) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("capture_metadata", [False, True]) 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 index 7ab9d301f8..cfbee30cc1 100644 --- 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 @@ -431,3 +431,40 @@ def step(): atol=0.002, rtol=torch.finfo(dtype).eps, ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_full_prefix_does_not_address_past_capacity(kda): + torch.manual_seed(4) + capacity = 4 + backend = ReplaySSMExecutor( + torch.randn(1, 3, 2, 32, 32, device="cuda") * 0.01, + capacity, + 1, + torch.bfloat16, + kda=kda, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + q = torch.randn(2, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(2, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + backend.ssm_state.zero_() + backend.cursors.zero_() + for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): + getattr(backend, name).zero_() + backend.cursors[reqs.long()] = capacity + 2 + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + backend.merge_accepted_updates(reqs) + torch.cuda.synchronize() + assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) + assert torch.isfinite(backend.ssm_state).all() + for _ in range(12): + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + backend.accept_single_token(reqs) + backend.maybe_compress_prefix(reqs) + torch.cuda.synchronize() + assert int(backend.cursors[reqs.long()].max()) <= capacity + assert torch.isfinite(backend.ssm_state).all() From fea39d6e6d775b8f829ef592591de45a939dbf9d Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 04:17:03 +0000 Subject: [PATCH 29/54] fix: fold replay prefixes inside the decode kernel The separate fold cleared the shared cursor while other layers were still reading it, so accepted history never reached the SSM and GSM8K trajectories collapsed. Record the write base during replay and advance the cursor only after every layer has finished. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 44 +++++-- .../bounded_att/kda/replay_kda_state.py | 44 +++++-- .../attention/linear/replay_gdn_decode.py | 1 + .../attention/linear/replay_kda_decode.py | 1 + .../triton_kernel/linear_att/replayssm.py | 3 +- .../linear_att/replayssm_cursor.py | 28 +++- .../linear_att/replayssm_decode.py | 31 ++++- .../linear_att/replayssm_fold.py | 34 ++++- .../linear_att/ssm_acceptance.py | 20 ++- test.md | 58 +++++++++ .../linear_att/test_acceptance_fusion.py | 6 + .../linear_att/test_replayssm.py | 120 +++++++++++++++++- .../linear_att/test_replayssm_execution.py | 12 +- .../common/test_linear_state_contract.py | 1 + 14 files changed, 345 insertions(+), 58 deletions(-) create mode 100644 test.md diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index b4f03b2602..4eb1cd8d1f 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -15,6 +15,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _advance_cursor, + _clear_cursor_if_full, _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix @@ -92,7 +93,6 @@ def __init__( replay_config = dict(ssm_replay_config or {}) capacity = int(replay_config.get("cache_len", 16)) compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - force_compress = bool(replay_config.get("force_compress", False)) _bind_history( self, self.ssm_state, @@ -103,7 +103,6 @@ def __init__( gate_lower_bound, compute_mode, None, - force_compress, ) def get_state_type(self) -> str: @@ -214,11 +213,10 @@ def restore_bounded_att_runtime_from_pages( self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + """提交已接受长度。前缀放不下时，折回发生在下一次 replay kernel 里。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - self.maybe_compress_prefix(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -418,13 +416,12 @@ def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: self.compress_prefix(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1]: return self.accept_single_token(reqs) - self.maybe_compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -432,8 +429,10 @@ def clear_history(self, req: int | torch.Tensor) -> None: if type(req) is int: # 标量赋值会把 CPU tensor 复制到 CUDA slot。 self.cursors[req].zero_() + self.write_bases[req].zero_() else: self.cursors[req] = 0 + self.write_bases[req] = 0 def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" @@ -448,6 +447,7 @@ def accept_updates( _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, + self.write_bases, reqs, accepted_state_indices, reqs.numel(), @@ -483,14 +483,21 @@ def compress_prefix(self, reqs: torch.Tensor) -> None: self.clear_history(reqs) def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: - """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + """长度加上一次最大 verify_width 放不下时，整段压缩，并在折回结束后清掉这些请求的 cursor。""" if reqs.numel() == 0 or not self.ssm_state.is_cuda: return - if self.force_compress: - self.compress_prefix(reqs) - return self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + _clear_cursor_if_full[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + self.write_bases, + reqs, + reqs.numel(), + self.hold_req_id, + self.capacity, + self.verify_width, + 256, + ) def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -525,12 +532,16 @@ def export_accepted_state( flush_export=flush, output_strides=strides, ) + if flush: + slots = torch.tensor(req_indices, dtype=torch.int64, device=device) + self.cursors.index_fill_(0, slots, 0) + self.write_bases.index_fill_(0, slots, 0) def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + """把接受长度写入 MTP 槽位索引，并把 cursor 设为这次 replay 的写起点加接受长度。""" if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) + commit_ssm_prefix(indices, reqs, lengths, self.cursors, self.write_bases) else: commit_ssm_prefix(indices, reqs, lengths) self.accept_updates(reqs, indices) @@ -711,6 +722,7 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape + layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} bv = config["BV"] _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( @@ -737,6 +749,11 @@ def _materialize_accepted( OUTPUT_STRIDES=output_strides, WHEN_FULL=when_full, NEXT=self.verify_width, + STRIDE_LAYER=layer_stride, + STRIDE_SLOT=slot_stride, + STRIDE_HEAD=head_stride, + STRIDE_K=key_stride, + STRIDE_V=value_stride, num_warps=config["num_warps"], num_stages=config["num_stages"], ) @@ -789,7 +806,6 @@ def _bind_history( gate_lower_bound: float, compute_mode: str, run_config: dict[str, int] | None, - force_compress: bool = False, ) -> None: """在 owner 上创建 GDN replay 的连续前缀 history 和 cursor。""" @@ -801,7 +817,6 @@ def _bind_history( owner.gate_lower_bound = gate_lower_bound owner.capacity = capacity owner.verify_width = verify_width - owner.force_compress = bool(force_compress) owner.compute_mode = compute_mode owner.run_config = run_config owner._config_is_fixed = run_config is not None @@ -811,6 +826,7 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 7df2df8738..237030e77d 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -15,6 +15,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _advance_cursor, + _clear_cursor_if_full, _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix @@ -92,7 +93,6 @@ def __init__( replay_config = dict(ssm_replay_config or {}) capacity = int(replay_config.get("cache_len", 16)) compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - force_compress = bool(replay_config.get("force_compress", False)) _bind_history( self, self.ssm_state, @@ -103,7 +103,6 @@ def __init__( gate_lower_bound, compute_mode, None, - force_compress, ) def get_state_type(self) -> str: @@ -214,11 +213,10 @@ def restore_bounded_att_runtime_from_pages( self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: - """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + """提交已接受长度。前缀放不下时，折回发生在下一次 replay kernel 里。""" if self.mtp_state_index is None: return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - self.maybe_compress_prefix(req_indices) def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" @@ -417,13 +415,12 @@ def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: self.compress_prefix(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1]: return self.accept_single_token(reqs) - self.maybe_compress_prefix(reqs) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -431,8 +428,10 @@ def clear_history(self, req: int | torch.Tensor) -> None: if type(req) is int: # 标量赋值会把 CPU tensor 复制到 CUDA slot。 self.cursors[req].zero_() + self.write_bases[req].zero_() else: self.cursors[req] = 0 + self.write_bases[req] = 0 def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" @@ -447,6 +446,7 @@ def accept_updates( _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, + self.write_bases, reqs, accepted_state_indices, reqs.numel(), @@ -482,14 +482,21 @@ def compress_prefix(self, reqs: torch.Tensor) -> None: self.clear_history(reqs) def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: - """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + """长度加上一次最大 verify_width 放不下时，整段压缩，并在折回结束后清掉这些请求的 cursor。""" if reqs.numel() == 0 or not self.ssm_state.is_cuda: return - if self.force_compress: - self.compress_prefix(reqs) - return self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + _clear_cursor_if_full[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + self.write_bases, + reqs, + reqs.numel(), + self.hold_req_id, + self.capacity, + self.verify_width, + 256, + ) def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -524,12 +531,16 @@ def export_accepted_state( flush_export=flush, output_strides=strides, ) + if flush: + slots = torch.tensor(req_indices, dtype=torch.int64, device=device) + self.cursors.index_fill_(0, slots, 0) + self.write_bases.index_fill_(0, slots, 0) def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + """把接受长度写入 MTP 槽位索引，并把 cursor 设为这次 replay 的写起点加接受长度。""" if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors) + commit_ssm_prefix(indices, reqs, lengths, self.cursors, self.write_bases) else: commit_ssm_prefix(indices, reqs, lengths) self.accept_updates(reqs, indices) @@ -710,6 +721,7 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape + layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} bv = config["BV"] _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( @@ -736,6 +748,11 @@ def _materialize_accepted( OUTPUT_STRIDES=output_strides, WHEN_FULL=when_full, NEXT=self.verify_width, + STRIDE_LAYER=layer_stride, + STRIDE_SLOT=slot_stride, + STRIDE_HEAD=head_stride, + STRIDE_K=key_stride, + STRIDE_V=value_stride, num_warps=config["num_warps"], num_stages=config["num_stages"], ) @@ -788,7 +805,6 @@ def _bind_history( gate_lower_bound: float, compute_mode: str, run_config: dict[str, int] | None, - force_compress: bool = False, ) -> None: """在 owner 上创建 KDA replay 的连续前缀 history 和 cursor。""" @@ -800,7 +816,6 @@ def _bind_history( owner.gate_lower_bound = gate_lower_bound owner.capacity = capacity owner.verify_width = verify_width - owner.force_compress = bool(force_compress) owner.compute_mode = compute_mode owner.run_config = run_config owner._config_is_fixed = run_config is not None @@ -810,6 +825,7 @@ def _bind_history( assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 71fcf8e1aa..37551d45b2 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -57,6 +57,7 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, + history.write_bases, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 55514aeb57..1694751f38 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -57,6 +57,7 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, + history.write_bases, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index e54186eae4..cbfe51bbf4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -14,7 +14,7 @@ _replay_each_token, ) from .replayssm_delta import _delta_rule_step -from .replayssm_cursor import _advance_cursor +from .replayssm_cursor import _advance_cursor, _clear_cursor_if_full from .replayssm_fold import _fold_prefix __all__ = [ @@ -25,5 +25,6 @@ "_replay_each_token", "_delta_rule_step", "_advance_cursor", + "_clear_cursor_if_full", "_fold_prefix", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py index 5354c71260..df3072e4e8 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -8,6 +8,7 @@ @triton.jit(do_not_specialize=["N"]) def _advance_cursor( cursors, + write_bases, reqs, accepted_state_indices, N, @@ -15,10 +16,31 @@ def _advance_cursor( WIDTH: tl.constexpr, BLOCK: tl.constexpr, ): - # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 + # 写起点由 replay kernel 记下。半区已经折回时起点是 0，否则是原 cursor。 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) + base = tl.load(write_bases + 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) + tl.store(cursors + r, base + count, valid) + + +@triton.jit(do_not_specialize=["N"]) +def _clear_cursor_if_full( + cursors, + write_bases, + reqs, + N, + HOLD: tl.constexpr, + CAPACITY: tl.constexpr, + NEXT: tl.constexpr, + BLOCK: tl.constexpr, +): + # 折回 kernel 已经结束，这里再把放不下下一次窗口的请求长度清零。 + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + n = tl.load(cursors + r, valid, 0) + full = valid & (n != 0) & (n + NEXT > CAPACITY) + tl.store(cursors + r, 0, full) + tl.store(write_bases + r, 0, full) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 5619869b64..1397af0a04 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -9,6 +9,8 @@ import triton import triton.language as tl +from .replayssm_delta import _delta_rule_step + @triton.jit def _history_coeff(rows, vec, BK: tl.constexpr): @@ -44,6 +46,7 @@ def _decode_replay( Betas, Reqs, Cursors, + WriteBase, Cu, Out, SQ: tl.constexpr, @@ -70,9 +73,8 @@ def _decode_replay( ): """执行一个请求窗口的 replay decode，并写入 history。 - 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 - cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， - 不应因为 eager batch 改变而形成新的 kernel 变体。 + 前缀放不下本步 token 时，先在本 kernel 里把已有前缀从原始 k/v 折进 checkpoint， + 舍入成 SSM 的 dtype，再从这份状态重放本步。cursor 本 kernel 不修改，写起点记到 WriteBase。 """ iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) req = tl.load(Reqs + seq).to(tl.int64) @@ -85,16 +87,37 @@ def _decode_replay( for t in range(start, end): tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) return + n = tl.load(Cursors + req) if start == end: + # 本步没有 token。写起点保持原 cursor，接受 0 个 token 时长度不变。 + if iv == 0 and hv == 0: + tl.store(WriteBase + req, n) return kk = tl.arange(0, BK) ll = tl.arange(0, L) slot = req * HV + hv key_slot = req * KH + hv // (HV // KH) - n = tl.load(Cursors + req) sp = State + slot * K * V + kk[:, None] * V + vv[None, :] state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + # 和半区方案一样：放不下时先折回，再在同一 kernel 里重放新 token。各层只读 cursor。 + if n + (end - start) > L: + for j in range(L): + active = j < n + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) + state = state.to(State.dtype.element_ty).to(tl.float32) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + n = 0 + if iv == 0 and hv == 0: + tl.store(WriteBase + req, n) if PRECOMPUTE_PROJECTIONS: # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 1a10d122dd..a01b9e6f1f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -32,10 +32,15 @@ def _fold_prefix( OUTPUT_STRIDES: tl.constexpr = None, WHEN_FULL: tl.constexpr = False, NEXT: tl.constexpr = 1, + STRIDE_LAYER: tl.constexpr = 0, + STRIDE_SLOT: tl.constexpr = 0, + STRIDE_HEAD: tl.constexpr = 0, + STRIDE_K: tl.constexpr = 0, + STRIDE_V: tl.constexpr = 1, ): # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT - # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， - # 并在 kernel 内把这些请求的 cursor 清零，避免主机分支进不了 CUDA Graph。 + # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求。 + # cursor 清零由随后的 kernel 完成，避免和还在读 cursor 的线程块竞争。 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: @@ -47,11 +52,25 @@ def _fold_prefix( elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: return linear_layer_idx, head = lh // HV, lh % HV + # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 + layer64 = linear_layer_idx.to(tl.int64) + head64 = head.to(tl.int64) + stride_layer: tl.constexpr = STRIDE_LAYER if STRIDE_LAYER != 0 else SLOTS * HV * K * V + stride_slot: tl.constexpr = STRIDE_SLOT if STRIDE_SLOT != 0 else HV * K * V + stride_head: tl.constexpr = STRIDE_HEAD if STRIDE_HEAD != 0 else K * V + stride_k: tl.constexpr = STRIDE_K if STRIDE_K != 0 else V 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, :] + sp = ( + State + + layer64 * stride_layer + + req * stride_slot + + head64 * stride_head + + kk[:, None].to(tl.int64) * stride_k + + vv[None, :].to(tl.int64) * STRIDE_V + ) state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 for j in range(L): @@ -75,11 +94,12 @@ def _fold_prefix( + kk[:, None] * OUTPUT_STRIDES[2] + vv[None, :] * OUTPUT_STRIDES[3] ) + elif SNAPSHOT: + out_slot = layer64 * HV + head64 + out = Output + out_slot * K * V + kk[:, None].to(tl.int64) * V + vv[None, :].to(tl.int64) else: - out_slot = linear_layer_idx * HV + head if SNAPSHOT else slot - out = Output + out_slot * K * V + kk[:, None] * V + vv[None, :] + out = sp tl.store(out, state, mask) if FLUSH_EXPORT: tl.store(sp, state, mask) - if (FLUSH_EXPORT or WHEN_FULL) and iv == 0 and lh == 0: - tl.store(Cursors + req, 0) + # cursor 不在这里清。和折回并行的其他线程块还要读原来的长度，清零放到 kernel 结束之后。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py index 5244dada83..86ad1c8025 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -10,13 +10,21 @@ def commit_ssm_prefix( reqs: torch.Tensor, lengths: torch.Tensor, replay_cursors: torch.Tensor | None = None, + write_bases: torch.Tensor | None = None, ) -> None: - """把每个请求的 accepted length 写入 MTP 索引并推进 Replay cursor。""" + """把每个请求的 accepted length 写入 MTP 索引。 + + 传入 write_bases 时，cursor 设为这次 replay 的写起点加接受长度。层内已经折回时写起点是 0。 + """ if not reqs.is_cuda: indices[reqs] = (lengths - 1).to(indices.dtype) if replay_cursors is not None: valid = reqs != replay_cursors.numel() - 1 - replay_cursors[reqs[valid]] += lengths[valid].to(replay_cursors.dtype) + accepted = lengths[valid].to(replay_cursors.dtype) + if write_bases is not None: + replay_cursors[reqs[valid]] = write_bases[reqs[valid]] + accepted + else: + replay_cursors[reqs[valid]] += accepted return if reqs.numel(): _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( @@ -24,11 +32,13 @@ def commit_ssm_prefix( reqs, lengths, replay_cursors, + write_bases if write_bases is not None else (replay_cursors if replay_cursors is not None else indices), reqs.numel(), indices.numel() - 1, reqs.stride(0), lengths.stride(0), 256, + WRITE_BASE=write_bases is not None, num_warps=4, ) @@ -39,11 +49,13 @@ def _commit_prefix( Reqs, Lengths, Cursors, + WriteBase, N, HOLD: tl.constexpr, REQ_STRIDE: tl.constexpr, LENGTH_STRIDE: tl.constexpr, BLOCK: tl.constexpr, + WRITE_BASE: tl.constexpr, ): row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) @@ -51,5 +63,5 @@ def _commit_prefix( tl.store(Indices + req, length - 1, row < N) if Cursors is not None: valid = (row < N) & (req != HOLD) - cursor = tl.load(Cursors + req, valid, 0) - tl.store(Cursors + req, cursor + length, valid) + base = tl.load(WriteBase + req, valid, 0) if WRITE_BASE else tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, base + length, valid) diff --git a/test.md b/test.md new file mode 100644 index 0000000000..f1a6298685 --- /dev/null +++ b/test.md @@ -0,0 +1,58 @@ +# ReplaySSM 折回：精度确认与两种时机的对比 + +测量对象是当前工作区的实现：普通 decode 放不下时，折回发生在每一层正在算线性注意力的 `replay` kernel 开头。容量够时只重放，不写 canonical SSM。Qwen3.5-0.8B，思考关闭，GSM8K 1000 题，`max_tokens=512`。 + +## 1. 1000 题上，现在这条 kernel 没有原先的精度崩溃 + +崩溃的样子是答案尾部刷出长串 `!` 后被 512 token 截断。同一套协议下，现在三条 replay 配置的 `!` 循环都是 0，strict 正确数贴着 native fp32。 + +| 配置 | strict | flexible | 截断 | `!` 循环 | 与 native 最终数字一致 | +|---|---:|---:|---:|---:|---:| +| native fp32 | 519 | 542 | 74 | 0 | 1000 | +| inline bf16，cache 4 | 511 | 534 | 81 | 0 | 807 | +| inline bf16，cache 16 | 515 | 536 | 63 | 0 | 810 | +| inline fp32，cache 16 | 518 | 533 | 67 | 0 | 803 | + +同一批题上，改层内折回之前： + +| 配置 | strict | 截断 | `!` 循环 | +|---|---:|---:|---:| +| inline bf16，cache 4 | 336 | 403 | 383 | +| inline bf16，cache 16 | 444 | 134 | 60 | +| inline fp32，cache 16 | 481 | 126 | 58 | + +和 native 对不上的大约两成，是回放和 fused kernel 的正常差别。cache 4 的正确数从 336 回到 511，已经贴着 native 的 519。这次没有重跑 precompute 和 KDA。 + +## 2. 整步结束后再折回，和层内折回的差别 + +「整步后再折回」指旧路径：每层 replay 先把 token 算完，最后一层 `maybe_compress_prefix` 再用独立 kernel `_fold_prefix` 折回，并在该 kernel 内把 cursor 写成 0。「层内折回」指现在的路径：放不下时，折回在每一层的 replay kernel 开头完成，共享 cursor 等全部层结束才在外面更新。 + +### 精度 + +独立折回如果等全部线程块写完再改 cursor，和层内折回是同一个结果。bf16 和 fp32、cache 4 和 cache 16，一步边界 token 的输出和 SSM 最大误差都是 0。 + +旧路径不是这样。`WHEN_FULL` 由其中一个线程块在 kernel 内部把 cursor 写成 0，其他线程块可能还没读到原来的长度，对应的层或头就跳过折回。和正确折回相比，状态最大误差： + +| 批量 | cache | 20 次里对不上 | 状态最大误差 | +|---|---:|---:|---:| +| 1 | 4 | 1 | 0.125 | +| 1 | 16 | 11 | 0.132 | +| 16 | 4 | 20 | 0.218 | +| 16 | 16 | 20 | 0.268 | + +这个误差会留在 SSM 里。cache 4 的 1000 题里有 530 题最终数字因此改变，其中 379 题改到与 native 一致，44 题离开 native。 + +### 性能 + +形状是 Qwen3.5-0.8B 的 18 个线性层、16 头、维度 128、bf16，CUDA graph。时间是 18 层加在一起。普通步是前缀已满但还没触发折回；折回步是刚好放不下的那一步。 + +| 批量 | cache | 普通回放 | 层内折回再回放 | 整步后的折回 kernel | 每步都有的空折回 | +|---|---:|---:|---:|---:|---:| +| 1 | 4 | 106 µs | 178 µs | 12 µs | 3.3 µs | +| 1 | 16 | 140 µs | 361 µs | 31 µs | 3.4 µs | +| 16 | 4 | 180 µs | 309 µs | 161 µs | 13 µs | +| 16 | 16 | 283 µs | 732 µs | 400 µs | 13 µs | + +单条请求时，折回串在每一层的 replay 里。cache 4 的那一步大约多 70 µs；独立 kernel 把 18 层一起折完只要 12 µs。摊到每个 token 大约多 10 µs。批量 16 时，独立折回本身要 160–400 µs，放进已经要启动的层 kernel 里，每个 token 大约少 10 µs。 + +旧路径在容量够的时候仍会每步启动一次空折回，批量 1 约 3 µs，批量 16 约 13 µs。现在这条路径不再发这个 kernel。相对完整 decode（还有 full attention 和 MLP），这几十微秒不是端到端的大头。1000 题墙钟时间现在是 55–61 秒，native 是 54 秒。 diff --git a/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 index fb5168f661..d03db1b99b 100644 --- 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 @@ -40,6 +40,7 @@ def make(n, width): args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] ssm_backend = ReplayGDNState.__new__(ReplayGDNState) ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.write_bases = ssm_backend.cursors.clone() ssm_backend.hold_req_id = hold ssm_backend.capacity = max(4, triton.next_power_of_2(width)) ssm_backend.verify_width = width @@ -113,6 +114,8 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): ref.mtp_state_index.fill_(-9) m._ssm_executor.cursors.fill_(7) ref._ssm_executor.cursors.fill_(7) + m._ssm_executor.write_bases.fill_(7) + ref._ssm_executor.write_bases.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): @@ -122,6 +125,9 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): if r != m._ssm_executor.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) + # replay 没有折回时，写起点就是当前 cursor。提交把它加上本步接受数。 + m._ssm_executor.write_bases.copy_(m._ssm_executor.cursors) + ref._ssm_executor.write_bases.copy_(ref._ssm_executor.cursors) accept_unfused(ref, args) if step % 2: graph.replay() 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 index 23270bb407..32829288d7 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -442,14 +442,87 @@ def check(capacity, width, steps): check(16, 2, 4) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capacity", [4, 8]) +def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): + # 放不下时，GDN 和 KDA 都在 replay kernel 内折回。cursor 等两层都算完才改。 + torch.manual_seed(11 + capacity + int(kda)) + hv, kh, dim = 2, 2, 32 + state = torch.randn(2, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) + twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) + assert backend.gates.shape[-1] == (dim if kda else capacity) + assert twin.write_bases.shape == backend.cursors.shape + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + for step in range(capacity + 1): + q = torch.randn(1, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(1, hv, dim, device="cuda", dtype=dtype) + a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() + b = torch.randn(1, hv, device="cuda", dtype=dtype) + if step == capacity: + twin.merge_accepted_updates(reqs) + assert int(backend.cursors[0]) == capacity + outs, twin_outs = [], [] + for layer in range(2): + outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs)[0]) + twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs)[0]) + if step == capacity: + assert int(backend.cursors[0]) == capacity + assert int(backend.write_bases[0]) == 0 + for actual, expected in zip(outs, twin_outs): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + backend.accept_single_token(reqs) + twin.accept_single_token(reqs) + assert int(backend.cursors[0]) == 1 + assert int(twin.cursors[0]) == 1 + # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 + state_atol = 0 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(backend.ssm_state, twin.ssm_state, atol=state_atol, rtol=0) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): + state = torch.randn(1, 3, 2, 32, 32, device="cuda") + backend = ReplaySSMExecutor( + state, 8, 4, torch.float32, num_key_heads=2, kda=kda, run_config={"BV": 32, "num_warps": 1} + ) + backend.cursors[0] = 3 + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) + q = torch.randn(1, 1, 2, 32, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 1, 2, 32, device="cuda") + a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") + b = torch.zeros(1, 2, device="cuda") + bias = torch.zeros(2 * 32 if kda else 2, device="cuda") + replay_layer(backend, 0, q, k, v, a, b, bias, bias, reqs, cu) + assert int(backend.write_bases[0]) == 3 + accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) + backend.accept_updates(reqs, accepted) + assert backend.cursors.tolist() == [3, 0, 0] + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("capture_metadata", [False, True]) -def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, capture_metadata): +def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_mode, capture_metadata): torch.manual_seed(31) state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) ssm_backend = ReplaySSMExecutor( - state, 8, 4, dtype, num_key_heads=2, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + state, + 8, + 4, + dtype, + num_key_heads=2, + kda=kda, + compute_mode=compute_mode, + run_config={"BV": 32, "num_warps": 1}, ) expected = state.float().clone() pending = [0, 0] @@ -459,12 +532,13 @@ def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, 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) + a = torch.full((8, 4 * 128 if kda else 4), -3.0, dtype=dtype, device="cuda") + b = torch.randn(8, 4, dtype=dtype, device="cuda") alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") def forward(): - return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] def step(): positions = ssm_backend.prepare_decode(reqs, cu) @@ -513,7 +587,7 @@ def step(): 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) + out, current = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, current, kda=kda) 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) @@ -725,3 +799,37 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa backend.merge_accepted_updates(reqs[:live]) torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) assert torch.equal(backend.ssm_state[0, live], initial[0, live]) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): + # 验证窗口放不下时，replay 先折回并把写起点记为 0。提交必须用写起点，不能再把旧 cursor 加上去。 + torch.manual_seed(19 + int(kda)) + hv, kh, dim = 2, 2, 32 + dtype = torch.bfloat16 + state = torch.randn(1, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + backend = ReplaySSMExecutor( + state, 4, 2, dtype, num_key_heads=kh, kda=kda, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + ) + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + indices = torch.zeros(2, device="cuda", dtype=torch.int32) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + cu = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + for step in range(3): + q = torch.randn(1, 2, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(1, 2, hv, dim, device="cuda", dtype=dtype) + a = torch.full((2, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) + b = torch.randn(2, hv, device="cuda", dtype=dtype) + replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) + lengths = torch.tensor([2], device="cuda", dtype=torch.int32) + backend.commit_accepted(indices, reqs, lengths) + if step < 2: + assert int(backend.cursors[0]) == 2 * (step + 1) + assert int(backend.write_bases[0]) == 2 * step + else: + assert int(backend.write_bases[0]) == 0 + assert int(backend.cursors[0]) == 2 + assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index cfbee30cc1..6c096165d3 100644 --- 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 @@ -319,7 +319,8 @@ def warm(ssm_backend): ) -def test_cross_layer_fold_with_state_offset_above_int32(): +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_cross_layer_fold_with_state_offset_above_int32(kda): if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: pytest.skip("large-stride regression requires 6 GiB of free device memory") torch.manual_seed(791) @@ -331,15 +332,16 @@ def test_cross_layer_fold_with_state_offset_above_int32(): sparse.copy_(dense) config = {"BV": 32, "num_warps": 1, "num_stages": 1} ssm_backends = [ - ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda, 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) + a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) alog = torch.zeros(12, device="cuda") - bias = torch.zeros_like(alog) + bias = torch.zeros(12 * 128 if kda else 12, device="cuda") reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) cu = reqs.new_tensor([0, 4, 8]) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index f720e01f09..a4df60ae25 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -194,6 +194,7 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): owner.mtp_state_index.fill_(2) history = getattr(owner, "runtime", owner) history.cursors.fill_(8) + history.write_bases.copy_(history.cursors) reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] owner.commit_accepted(owner.mtp_state_index, reqs, lengths) From c63ad74784c32be8c5e25f142d1f59b6fbe9f7ef Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 06:07:42 +0000 Subject: [PATCH 30/54] refactor: share the accepted-prefix fold between replay and compress Decode and the separate fold kernel each walked the same history loop. One device function keeps those two paths on the same update. Co-authored-by: Cursor --- .../triton_kernel/linear_att/replayssm.py | 3 +- .../linear_att/replayssm_decode.py | 20 +++--------- .../linear_att/replayssm_delta.py | 32 +++++++++++++++++++ .../linear_att/replayssm_fold.py | 14 ++------ 4 files changed, 41 insertions(+), 28 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index cbfe51bbf4..8551212f39 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -13,7 +13,7 @@ _project_qk, _replay_each_token, ) -from .replayssm_delta import _delta_rule_step +from .replayssm_delta import _apply_accepted_prefix, _delta_rule_step from .replayssm_cursor import _advance_cursor, _clear_cursor_if_full from .replayssm_fold import _fold_prefix @@ -23,6 +23,7 @@ "_gdn_precomputed_window", "_kda_precomputed_window", "_replay_each_token", + "_apply_accepted_prefix", "_delta_rule_step", "_advance_cursor", "_clear_cursor_if_full", diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 1397af0a04..df76b0ab20 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 """ReplaySSM decode kernels. -`_decode_replay` is the entry. It projects the verify window when asked, then -solves a GDN or KDA window, or walks tokens one by one. Neither path writes the -canonical SSM. +`_decode_replay` is the entry. When the accepted prefix does not fit, it folds +that prefix into the SSM and then replays the new tokens from the rounded +state. The output is either a precomputed GDN or KDA window, or a token walk. """ import triton import triton.language as tl -from .replayssm_delta import _delta_rule_step +from .replayssm_delta import _apply_accepted_prefix @triton.jit @@ -102,17 +102,7 @@ def _decode_replay( state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) # 和半区方案一样：放不下时先折回，再在同一 kernel 里重放新 token。各层只读 cursor。 if n + (end - start) > L: - for j in range(L): - active = j < n - raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) - raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) - if KDA: - g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) - else: - g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) - beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) - updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) - state = tl.where(active, updated, state) + state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) n = 0 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py index b3a6b2aec0..468246bf18 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -19,3 +19,35 @@ def _delta_rule_step(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): d = beta * (raw_v - decay * sk) state = state * decay + raw_k[:, None] * d[None, :] return state + + +@triton.jit +def _apply_accepted_prefix( + state, + RawKeys, + RawValues, + Gates, + Betas, + key_slot, + slot, + kk, + vv, + n, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + KDA: tl.constexpr, +): + # 按 history 槽位顺序把已接受前缀折进 state。j >= n 的槽不读、不更新。 + for j in range(L): + active = j < n + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) + return state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index a01b9e6f1f..a569d483d1 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from .replayssm_delta import _delta_rule_step +from .replayssm_delta import _apply_accepted_prefix @triton.jit @@ -73,17 +73,7 @@ def _fold_prefix( ) state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 - for j in range(L): - active = j < n - raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) - raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) - if KDA: - g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) - else: - g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) - beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) - updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) - state = tl.where(active, updated, state) + state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) mask = (kk[:, None] < K) & (vv[None, :] < V) if OUTPUT_STRIDES is not None: target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) From 244c9866014c9a6ca2940b1d5d658da3c935d58c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 07:16:02 +0000 Subject: [PATCH 31/54] refactor: remove the unused full-prefix replay path Acceptance and the decode kernel already keep the single prefix cursor. History setup stays in the constructor, and tests attach an existing SSM themselves. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 177 +++++------------- .../bounded_att/kda/replay_kda_state.py | 177 +++++------------- .../triton_kernel/linear_att/replayssm.py | 3 +- .../linear_att/replayssm_cursor.py | 21 --- .../linear_att/replayssm_fold.py | 11 +- test/benchmark/kernels/benchmark_kda_ssm.py | 15 +- .../linear_att/replay_history.py | 68 +++++++ .../linear_att/test_acceptance_fusion.py | 13 +- .../linear_att/test_replayssm.py | 20 +- .../linear_att/test_replayssm_execution.py | 22 +-- .../linear_att/test_ssm_autotune_memory.py | 13 +- 11 files changed, 182 insertions(+), 358 deletions(-) create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 4eb1cd8d1f..bf83af5c8e 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -15,7 +15,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _advance_cursor, - _clear_cursor_if_full, _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix @@ -92,18 +91,41 @@ def __init__( self.small_pages_pool = None replay_config = dict(ssm_replay_config or {}) capacity = int(replay_config.get("cache_len", 16)) + # 宽度为 1 逐 token 重放。开了 MTP 时默认先投影整个验证窗口。 compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - _bind_history( - self, - self.ssm_state, - capacity, - mtp_step + 1, - config.conv_dtype, - config.num_k_heads, - gate_lower_bound, - compute_mode, - None, - ) + verify_width = mtp_step + 1 + # 连续前缀挂在上面的 canonical SSM 上，不再另分配状态。 + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert self.ssm_state.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + self.kda = False + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + self.run_config = None + self._config_is_fixed = False + self.batch_configs = {} + layers, slots, hv, k, v = self.ssm_state.shape + self.num_key_heads = config.num_k_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 + self.hold_req_id = slots - 1 + # cursor 是已接受、尚未折回的长度，各层只读。 + # write_bases 是本步 replay 的写入起点：层内折回后为 0，否则等于进入 kernel 时的 cursor。 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + self.write_bases = torch.zeros(slots, dtype=torch.int32, device=device) + # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 + # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 + # GDN 的 g 是每个头一个标量。 + reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, self.num_key_heads, capacity, k) + self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=device) + self.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) + self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) + self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) + self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" @@ -412,16 +434,21 @@ def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层把已接受前缀整段压进 SSM。""" layers = () if self.config is None else self.config.logical_layer_ids - if layers and logical_layer_id == layers[-1]: - self.compress_prefix(reqs) + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1]: + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - self.accept_single_token(reqs) + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -456,14 +483,6 @@ def accept_updates( 256, ) - def accept_single_token(self, reqs: torch.Tensor) -> None: - """把一个已接受 token 记入前缀长度。""" - - if reqs.numel() == 0: - return - # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 - self.accept_updates(reqs, reqs, width=1) - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -472,33 +491,6 @@ def merge_accepted_updates(self, reqs: torch.Tensor) -> None: self._materialize_accepted(reqs, self.ssm_state, snapshot=False) self.cursors.index_fill_(0, reqs.long(), 0) - def compress_prefix(self, reqs: torch.Tensor) -> None: - """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" - - if reqs.numel() == 0: - return - if self.ssm_state.is_cuda: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) - - def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: - """长度加上一次最大 verify_width 放不下时，整段压缩，并在折回结束后清掉这些请求的 cursor。""" - - if reqs.numel() == 0 or not self.ssm_state.is_cuda: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) - _clear_cursor_if_full[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - self.write_bases, - reqs, - reqs.numel(), - self.hold_req_id, - self.capacity, - self.verify_width, - 256, - ) - def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -716,7 +708,6 @@ def _materialize_accepted( snapshot: bool, output_strides=None, flush_export: bool = False, - when_full: bool = False, ) -> None: """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" @@ -747,8 +738,6 @@ def _materialize_accepted( False, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, - WHEN_FULL=when_full, - NEXT=self.verify_width, STRIDE_LAYER=layer_stride, STRIDE_SLOT=slot_stride, STRIDE_HEAD=head_stride, @@ -757,83 +746,3 @@ def _materialize_accepted( num_warps=config["num_warps"], num_stages=config["num_stages"], ) - - -def _history_owner( - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - num_key_heads: int | None = None, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: dict[str, int] | None = None, -) -> ReplayGDNState: - """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" - - owner = ReplayGDNState.__new__(ReplayGDNState) - owner.config = None - owner.state_name = None - owner.request_capacity = ssm.shape[1] - owner.mtp_step = verify_width - 1 - owner._algorithm = "gdn" - owner._logical_to_storage_layer = {} - owner.conv_state = None - owner.mtp_state_index = None - owner.big_pages_pool = None - owner.small_pages_pool = None - _bind_history( - owner, - ssm, - capacity, - verify_width, - activation_dtype, - num_key_heads, - gate_lower_bound, - compute_mode, - run_config, - ) - return owner - - -def _bind_history( - owner, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype, - num_key_heads: int | None, - gate_lower_bound: float, - compute_mode: str, - run_config: dict[str, int] | None, -) -> None: - """在 owner 上创建 GDN replay 的连续前缀 history 和 cursor。""" - - assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - owner.ssm_state = ssm - owner.kda = False - owner.gate_lower_bound = gate_lower_bound - owner.capacity = capacity - owner.verify_width = verify_width - owner.compute_mode = compute_mode - owner.run_config = run_config - owner._config_is_fixed = run_config is not None - owner.batch_configs = {} - layers, slots, hv, k, v = ssm.shape - owner.num_key_heads = hv if num_key_heads is None else num_key_heads - assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 - owner.hold_req_id = slots - 1 - owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 - # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 - reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 - key_shape = (layers, slots, owner.num_key_heads, capacity, k) - owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) - owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 237030e77d..68b44b48c3 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -15,7 +15,6 @@ from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( _advance_cursor, - _clear_cursor_if_full, _fold_prefix, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix @@ -92,18 +91,41 @@ def __init__( self.small_pages_pool = None replay_config = dict(ssm_replay_config or {}) capacity = int(replay_config.get("cache_len", 16)) + # 宽度为 1 逐 token 重放。开了 MTP 时默认先投影整个验证窗口。 compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") - _bind_history( - self, - self.ssm_state, - capacity, - mtp_step + 1, - config.conv_dtype, - config.num_k_heads, - gate_lower_bound, - compute_mode, - None, - ) + verify_width = mtp_step + 1 + # 连续前缀挂在上面的 canonical SSM 上，不再另分配状态。 + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert self.ssm_state.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + self.kda = True + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + self.run_config = None + self._config_is_fixed = False + self.batch_configs = {} + layers, slots, hv, k, v = self.ssm_state.shape + self.num_key_heads = config.num_k_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 + self.hold_req_id = slots - 1 + # cursor 是已接受、尚未折回的长度，各层只读。 + # write_bases 是本步 replay 的写入起点：层内折回后为 0，否则等于进入 kernel 时的 cursor。 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + self.write_bases = torch.zeros(slots, dtype=torch.int32, device=device) + # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 + # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 + # KDA 的 g 按 key 通道存。 + reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, self.num_key_heads, capacity, k) + self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=device) + self.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=device) + self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) + self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) + self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) def get_state_type(self) -> str: """返回算法名和 Conv/SSM 存储精度。""" @@ -411,16 +433,21 @@ def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层把已接受前缀整段压进 SSM。""" layers = () if self.config is None else self.config.logical_layer_ids - if layers and logical_layer_id == layers[-1]: - self.compress_prefix(reqs) + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1]: + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - self.accept_single_token(reqs) + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) def clear_history(self, req: int | torch.Tensor) -> None: """清空复用或恢复请求的临时 history，保留 canonical SSM。""" @@ -455,14 +482,6 @@ def accept_updates( 256, ) - def accept_single_token(self, reqs: torch.Tensor) -> None: - """把一个已接受 token 记入前缀长度。""" - - if reqs.numel() == 0: - return - # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 - self.accept_updates(reqs, reqs, width=1) - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" @@ -471,33 +490,6 @@ def merge_accepted_updates(self, reqs: torch.Tensor) -> None: self._materialize_accepted(reqs, self.ssm_state, snapshot=False) self.cursors.index_fill_(0, reqs.long(), 0) - def compress_prefix(self, reqs: torch.Tensor) -> None: - """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" - - if reqs.numel() == 0: - return - if self.ssm_state.is_cuda: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) - - def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: - """长度加上一次最大 verify_width 放不下时，整段压缩，并在折回结束后清掉这些请求的 cursor。""" - - if reqs.numel() == 0 or not self.ssm_state.is_cuda: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) - _clear_cursor_if_full[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - self.write_bases, - reqs, - reqs.numel(), - self.hold_req_id, - self.capacity, - self.verify_width, - 256, - ) - def accepted_ssm(self, req_idx: int) -> torch.Tensor: """返回不修改 canonical SSM 和 history 的单请求快照。""" @@ -715,7 +707,6 @@ def _materialize_accepted( snapshot: bool, output_strides=None, flush_export: bool = False, - when_full: bool = False, ) -> None: """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" @@ -746,8 +737,6 @@ def _materialize_accepted( True, FLUSH_EXPORT=flush_export, OUTPUT_STRIDES=output_strides, - WHEN_FULL=when_full, - NEXT=self.verify_width, STRIDE_LAYER=layer_stride, STRIDE_SLOT=slot_stride, STRIDE_HEAD=head_stride, @@ -756,83 +745,3 @@ def _materialize_accepted( num_warps=config["num_warps"], num_stages=config["num_stages"], ) - - -def _history_owner( - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype = torch.bfloat16, - num_key_heads: int | None = None, - gate_lower_bound: float = -5.0, - compute_mode: str = "inline", - run_config: dict[str, int] | None = None, -) -> ReplayKDAState: - """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" - - owner = ReplayKDAState.__new__(ReplayKDAState) - owner.config = None - owner.state_name = None - owner.request_capacity = ssm.shape[1] - owner.mtp_step = verify_width - 1 - owner._algorithm = "kda" - owner._logical_to_storage_layer = {} - owner.conv_state = None - owner.mtp_state_index = None - owner.big_pages_pool = None - owner.small_pages_pool = None - _bind_history( - owner, - ssm, - capacity, - verify_width, - activation_dtype, - num_key_heads, - gate_lower_bound, - compute_mode, - run_config, - ) - return owner - - -def _bind_history( - owner, - ssm: torch.Tensor, - capacity: int, - verify_width: int, - activation_dtype: torch.dtype, - num_key_heads: int | None, - gate_lower_bound: float, - compute_mode: str, - run_config: dict[str, int] | None, -) -> None: - """在 owner 上创建 KDA replay 的连续前缀 history 和 cursor。""" - - assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute") - owner.ssm_state = ssm - owner.kda = True - owner.gate_lower_bound = gate_lower_bound - owner.capacity = capacity - owner.verify_width = verify_width - owner.compute_mode = compute_mode - owner.run_config = run_config - owner._config_is_fixed = run_config is not None - owner.batch_configs = {} - layers, slots, hv, k, v = ssm.shape - owner.num_key_heads = hv if num_key_heads is None else num_key_heads - assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 - owner.hold_req_id = slots - 1 - owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - # 输出重建用归一化 k 和 d。bf16 激活时这两份用 fp16，尾数比 bf16 多，和当前精度更好的重建缓存一致。 - # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 - reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 - key_shape = (layers, slots, owner.num_key_heads, capacity, k) - owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) - owner.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=ssm.device) - owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) - owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) - owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 8551212f39..c0a598143e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -14,7 +14,7 @@ _replay_each_token, ) from .replayssm_delta import _apply_accepted_prefix, _delta_rule_step -from .replayssm_cursor import _advance_cursor, _clear_cursor_if_full +from .replayssm_cursor import _advance_cursor from .replayssm_fold import _fold_prefix __all__ = [ @@ -26,6 +26,5 @@ "_apply_accepted_prefix", "_delta_rule_step", "_advance_cursor", - "_clear_cursor_if_full", "_fold_prefix", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py index df3072e4e8..5eee325083 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -23,24 +23,3 @@ def _advance_cursor( base = tl.load(write_bases + r, valid, 0) count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 tl.store(cursors + r, base + count, valid) - - -@triton.jit(do_not_specialize=["N"]) -def _clear_cursor_if_full( - cursors, - write_bases, - reqs, - N, - HOLD: tl.constexpr, - CAPACITY: tl.constexpr, - NEXT: tl.constexpr, - BLOCK: tl.constexpr, -): - # 折回 kernel 已经结束，这里再把放不下下一次窗口的请求长度清零。 - i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - r = tl.load(reqs + i, i < N, HOLD) - valid = (i < N) & (r != HOLD) - n = tl.load(cursors + r, valid, 0) - full = valid & (n != 0) & (n + NEXT > CAPACITY) - tl.store(cursors + r, 0, full) - tl.store(write_bases + r, 0, full) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index a569d483d1..e8742846aa 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -30,26 +30,19 @@ def _fold_prefix( KDA: tl.constexpr, FLUSH_EXPORT: tl.constexpr = False, OUTPUT_STRIDES: tl.constexpr = None, - WHEN_FULL: tl.constexpr = False, - NEXT: tl.constexpr = 1, STRIDE_LAYER: tl.constexpr = 0, STRIDE_SLOT: tl.constexpr = 0, STRIDE_HEAD: tl.constexpr = 0, STRIDE_K: tl.constexpr = 0, STRIDE_V: tl.constexpr = 1, ): - # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT - # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求。 - # cursor 清零由随后的 kernel 完成，避免和还在读 cursor 的线程块竞争。 + # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT 只写独立输出。 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 n = tl.load(Cursors + req) - if WHEN_FULL: - if n == 0 or n + NEXT <= L: - return - elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + if n == 0 and not SNAPSHOT and not FLUSH_EXPORT: return linear_layer_idx, head = lh // HV, lh % HV # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index f4928a9017..733fc36b50 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -14,19 +14,10 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda - - -def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.pop("kda", False) - if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - - return _history_owner(*args, **kwargs) - - from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) def workload(args, batch, width, dtype, mode, acceptance): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py new file mode 100644 index 0000000000..916291e26b --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -0,0 +1,68 @@ +"""测试用的 replay history。正式路径走 ReplayGDNState / ReplayKDAState 的构造函数。""" + +import torch + +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + + +def attach_replay_history( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, + *, + kda: bool = False, +): + """把连续前缀挂到调用方已有的 SSM 上，保留它原来的 stride。 + + kernel 测试需要非连续的层距，不能再走会重新分配状态的构造函数。 + """ + + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert ssm.dtype in (torch.float32, torch.bfloat16), "replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + state_cls = ReplayKDAState if kda else ReplayGDNState + owner = state_cls.__new__(state_cls) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "kda" if kda else "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + owner.ssm_state = ssm + owner.kda = kda + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度，write_bases 是本步写入起点。 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + # 重建用 k/d：bf16 激活用 fp16。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + gate_shape = (layers, slots, hv, capacity, k) if kda else (layers, slots, hv, capacity) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) + owner.gates = torch.empty(gate_shape, dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) + return owner 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 index d03db1b99b..2b111e41cc 100644 --- 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 @@ -3,16 +3,9 @@ import triton from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update - - -def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.pop("kda", False) - if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - - return _history_owner(*args, **kwargs) +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) def ssm_snapshot(backend, req): 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 index 32829288d7..1d836a77bc 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,15 +1,9 @@ import pytest import torch - -def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.pop("kda", False) - if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - - return _history_owner(*args, **kwargs) +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) def ssm_snapshot(backend, req): @@ -433,7 +427,9 @@ def check(capacity, width, steps): torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) for backend in backends: backend.accept_updates(reqs, accepted, width=width) - backend.maybe_compress_prefix(reqs) + overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] + if overflow.numel(): + backend.merge_accepted_updates(overflow) for backend in backends: backend.merge_accepted_updates(reqs) torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) @@ -476,8 +472,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): assert int(backend.write_bases[0]) == 0 for actual, expected in zip(outs, twin_outs): torch.testing.assert_close(actual, expected, atol=0, rtol=0) - backend.accept_single_token(reqs) - twin.accept_single_token(reqs) + backend.accept_updates(reqs, reqs, width=1) + twin.accept_updates(reqs, reqs, width=1) assert int(backend.cursors[0]) == 1 assert int(twin.cursors[0]) == 1 # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 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 index 6c096165d3..000810d22b 100644 --- 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 @@ -1,17 +1,9 @@ import pytest import torch - -def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.pop("kda", False) - if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - - return _history_owner(*args, **kwargs) - - +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, @@ -56,7 +48,7 @@ def cycle(backend, ids): count = ids.numel() backend.prepare_decode(ids) out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] - backend.accept_single_token(ids) + backend.accept_updates(ids, ids, width=1) return out graphs = {} @@ -465,8 +457,10 @@ def test_full_prefix_does_not_address_past_capacity(kda): assert torch.isfinite(backend.ssm_state).all() for _ in range(12): replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) - backend.accept_single_token(reqs) - backend.maybe_compress_prefix(reqs) + backend.accept_updates(reqs, reqs, width=1) + overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] + if overflow.numel(): + backend.merge_accepted_updates(overflow) torch.cuda.synchronize() assert int(backend.cursors[reqs.long()].max()) <= capacity assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index 726eec31c4..c6312f5f9a 100644 --- 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 @@ -2,16 +2,9 @@ import torch from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning - - -def ReplaySSMExecutor(*args, **kwargs): - kda = kwargs.pop("kda", False) - if kda: - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner - else: - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner - - return _history_owner(*args, **kwargs) +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") From 3a8540b1c162ba3d5ac96c283bb91cc102762e27 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 07:36:11 +0000 Subject: [PATCH 32/54] refactor: drop the unused replay prepare_decode hook Decode reads the prefix cursor inside the kernel. Tests that need the length read it directly. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 6 ------ .../bounded_att/kda/replay_kda_state.py | 6 ------ .../triton_kernel/linear_att/ssm_autotune.py | 2 +- test/benchmark/kernels/benchmark_kda_ssm.py | 2 -- .../linear_att/test_acceptance_fusion.py | 1 - .../triton_kernel/linear_att/test_replayssm.py | 8 ++------ .../linear_att/test_replayssm_execution.py | 11 ++++------- .../linear_att/test_ssm_autotune_memory.py | 1 - 8 files changed, 7 insertions(+), 30 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index bf83af5c8e..23da086d8f 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -461,12 +461,6 @@ def clear_history(self, req: int | torch.Tensor) -> None: self.cursors[req] = 0 self.write_bases[req] = 0 - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" - - del cu_seqlens - return self.cursors.index_select(0, reqs.long()) - def accept_updates( self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None ) -> None: diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 68b44b48c3..36d88460e1 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -460,12 +460,6 @@ def clear_history(self, req: int | torch.Tensor) -> None: self.cursors[req] = 0 self.write_bases[req] = 0 - def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: - """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" - - del cu_seqlens - return self.cursors.index_select(0, reqs.long()) - def accept_updates( self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None ) -> None: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 3a25fd5f9f..f9404e6114 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -160,7 +160,7 @@ def select_config( ssm_backend.run_config = run_config reqs, accepted = workload ssm_backend.cursors.zero_() - # 把前缀写满但不越出 capacity。replay 自己读 cursor，正式 decode 也不经过 prepare_decode。 + # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。 tokens = q.shape[1 if cu_seqlens is not None else 0] rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) for _ in range(rounds): diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 733fc36b50..c50adef255 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -55,8 +55,6 @@ def workload(args, batch, width, dtype, mode, acceptance): 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( 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 index 2b111e41cc..a826fb786b 100644 --- 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 @@ -182,7 +182,6 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput def cycle(i): ssm_backend = ssm_backends[i] - ssm_backend.prepare_decode(reqs, cu) outputs = [replay_layer(ssm_backend, 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 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 index 1d836a77bc..0d0f1ec417 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -64,7 +64,6 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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 history correction by reconstructing from raw inputs. @@ -174,7 +173,6 @@ def test_replay_acceptance_flush_and_graph( accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): - ssm_backend.prepare_decode(reqs, cu) out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] ssm_backend.accept_updates(reqs, accepted) return out @@ -347,7 +345,6 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d for ssm_backend in ssm_backends: ssm_backend.merge_accepted_updates(overflow_reqs) for ssm_backend in ssm_backends: - ssm_backend.prepare_decode(reqs, cu) outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) ssm_backend.accept_updates(reqs, accepted) if dtype == torch.float32: @@ -537,7 +534,7 @@ def forward(): return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] def step(): - positions = ssm_backend.prepare_decode(reqs, cu) + positions = ssm_backend.cursors.index_select(0, reqs.long()) out = forward() ssm_backend.accept_updates(reqs, accepted) return positions, out @@ -573,7 +570,7 @@ def step(): pending[req] = 0 old_cursors = [pending[req] for req in ids] + [0] if not capture_metadata: - positions.copy_(ssm_backend.prepare_decode(reqs, cu)) + positions.copy_(ssm_backend.cursors.index_select(0, reqs.long())) graph.replay() if not capture_metadata: ssm_backend.accept_updates(reqs, accepted) @@ -763,7 +760,6 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa b = torch.randn(live, query_width, 1, device="cuda") if (backend.cursors[:live] + query_width > backend.capacity).any(): backend.merge_accepted_updates(reqs[:live]) - backend.prepare_decode(reqs, cu) actual = replay_layer( backend, 0, 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 index 000810d22b..83fe267642 100644 --- 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 @@ -46,7 +46,6 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() - backend.prepare_decode(ids) out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] backend.accept_updates(ids, ids, width=1) return out @@ -146,7 +145,7 @@ def inputs(iteration): def metadata(ssm_backend): cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) - return cu, reqs, ssm_backend.prepare_decode(reqs, cu) + return cu, reqs, ssm_backend.cursors.index_select(0, reqs.long()) def forward_accept(index, cu, reqs): ssm_backend = ssm_backends[index] @@ -255,7 +254,6 @@ def forward(ssm_backend): return [replay_layer(ssm_backend, 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 @@ -274,7 +272,7 @@ def warm(ssm_backend): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): if capture_metadata: - pos = ssm_backend.prepare_decode(reqs) + pos = ssm_backend.cursors.index_select(0, reqs.long()) out = forward(ssm_backend) if capture_metadata: ssm_backend.accept_updates(reqs) @@ -293,7 +291,7 @@ def warm(ssm_backend): old = ssm_backends[0].cursors[reqs.long()].clone() for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): if not capture_metadata: - pos.copy_(ssm_backend.prepare_decode(reqs)) + pos.copy_(ssm_backend.cursors.index_select(0, reqs.long())) graph.replay() if not capture_metadata: ssm_backend.accept_updates(reqs) @@ -341,7 +339,7 @@ def step(ssm_backend, accepted): overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] if overflow: ssm_backend.merge_accepted_updates(reqs.new_tensor(overflow)) - pos = ssm_backend.prepare_decode(reqs, cu) + pos = ssm_backend.cursors.index_select(0, reqs.long()) out = [replay_layer(ssm_backend, 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 @@ -390,7 +388,6 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): def step(): - ssm_backend.prepare_decode(ids, offsets) out = [ replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(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 index c6312f5f9a..8f3f8ec976 100644 --- 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 @@ -70,6 +70,5 @@ def oom(*args, **kwargs): 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) tuning.replay_layer(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) torch.cuda.synchronize() From e73ae2e07b7dcf5c8388f156c5d3bb4fb87f0ec7 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 07:55:16 +0000 Subject: [PATCH 33/54] refactor: fill replay kernel defaults at launch Mode checks are asserts. The state stores the caller's config unchanged, and replay and fold supply missing BV, warps, and stages when the kernel starts. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 35 +++++-------------- .../bounded_att/kda/replay_kda_state.py | 35 +++++-------------- .../attention/linear/replay_gdn_decode.py | 12 ++++--- .../attention/linear/replay_kda_decode.py | 12 ++++--- .../triton_kernel/linear_att/ssm_autotune.py | 2 +- 5 files changed, 33 insertions(+), 63 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 23da086d8f..ca0e8e7d04 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -54,8 +54,7 @@ def __init__( gate_lower_bound: float = -5.0, ) -> None: """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" - if ssm_state_mode != "replay": - raise ValueError(f"ReplayGDNState only implements replay mode, got {ssm_state_mode}") + assert ssm_state_mode == "replay", f"ReplayGDNState only implements replay mode, got {ssm_state_mode}" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config self.state_name = None @@ -548,27 +547,6 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.ssm_state[:, req].copy_(snap) self.clear_history(req) - _run_config: dict[str, int] | None = None - - @property - def run_config(self) -> dict[str, int] | None: - """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: dict | None) -> None: - """保存一份与调用方脱开的 kernel 配置。""" - - if config is None: - self._run_config = None - return - self._run_config = { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } - # ======================================================================== # 七、私有辅助实现 # ======================================================================== @@ -708,8 +686,11 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() - config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} - bv = config["BV"] + # 折回不走 replay 的 autotune。没固定配置时用这份默认值。 + config = self.run_config or {} + bv = config.get("BV", 32) + num_warps = config.get("num_warps", 4) + num_stages = config.get("num_stages", 3) _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, @@ -737,6 +718,6 @@ def _materialize_accepted( STRIDE_HEAD=head_stride, STRIDE_K=key_stride, STRIDE_V=value_stride, - num_warps=config["num_warps"], - num_stages=config["num_stages"], + num_warps=num_warps, + num_stages=num_stages, ) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 36d88460e1..874acdb070 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -54,8 +54,7 @@ def __init__( gate_lower_bound: float = -5.0, ) -> None: """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" - if ssm_state_mode != "replay": - raise ValueError(f"ReplayKDAState only implements replay mode, got {ssm_state_mode}") + assert ssm_state_mode == "replay", f"ReplayKDAState only implements replay mode, got {ssm_state_mode}" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config self.state_name = None @@ -547,27 +546,6 @@ def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: self.ssm_state[:, req].copy_(snap) self.clear_history(req) - _run_config: dict[str, int] | None = None - - @property - def run_config(self) -> dict[str, int] | None: - """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" - - return self._run_config - - @run_config.setter - def run_config(self, config: dict | None) -> None: - """保存一份与调用方脱开的 kernel 配置。""" - - if config is None: - self._run_config = None - return - self._run_config = { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } - # ======================================================================== # 七、私有辅助实现 # ======================================================================== @@ -707,8 +685,11 @@ def _materialize_accepted( assert not (snapshot and flush_export), "flush export must update the canonical state" layers, slots, hv, k, v = self.ssm_state.shape layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() - config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} - bv = config["BV"] + # 折回不走 replay 的 autotune。没固定配置时用这份默认值。 + config = self.run_config or {} + bv = config.get("BV", 32) + num_warps = config.get("num_warps", 4) + num_stages = config.get("num_stages", 3) _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( self.ssm_state, output, @@ -736,6 +717,6 @@ def _materialize_accepted( STRIDE_HEAD=head_stride, STRIDE_K=key_stride, STRIDE_V=value_stride, - num_warps=config["num_warps"], - num_stages=config["num_stages"], + num_warps=num_warps, + num_stages=num_stages, ) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 37551d45b2..36b589563e 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -9,7 +9,7 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config def replay( @@ -28,9 +28,13 @@ def replay( ) -> tuple[torch.Tensor, dict[str, int]]: """执行一层 GDN replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" - if history.kda: - raise ValueError("GDN replay step does not use the KDA gate") - config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + assert not history.kda, "GDN replay step does not use the KDA gate" + # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 + config = ( + _replay_config(run_config) + if run_config is not None + else resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + ) hv, kd, vd = history.ssm_state.shape[-3:] assert history.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 1694751f38..de99c5fba4 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -9,7 +9,7 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config def replay( @@ -28,9 +28,13 @@ def replay( ) -> tuple[torch.Tensor, dict[str, int]]: """执行一层 KDA replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" - if not history.kda: - raise ValueError("KDA replay step requires the KDA gate") - config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + assert history.kda, "KDA replay step requires the KDA gate" + # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 + config = ( + _replay_config(run_config) + if run_config is not None + else resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + ) hv, kd, vd = history.ssm_state.shape[-3:] assert history.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index f9404e6114..2864f4f318 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -18,7 +18,7 @@ def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: """返回当前 batch 的固定或已调优 SSM kernel 配置。""" if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: - return ssm_backend.run_config + return _replay_config(ssm_backend.run_config) default = _replay_config(ssm_backend.run_config) key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) if key in ssm_backend.batch_configs: From 5b8623a7e18643b0df4e4190b6206048c03a8211 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 08:13:30 +0000 Subject: [PATCH 34/54] refactor: copy the prefill baseline into replay pages Page save, PD, and DP run before any replay history exists. They copy the accepted conv window and canonical SSM directly. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 49 +++++++------------ .../bounded_att/kda/replay_kda_state.py | 49 +++++++------------ .../linear_att/test_replayssm.py | 6 +-- .../test_hybrid_req_manager_pd_sharing.py | 9 ++-- unit_tests/common/test_linear_state.py | 9 ++-- unit_tests/models/glm5_next/test_cache.py | 1 - unit_tests/models/glm5_next/test_pd_cache.py | 1 - 7 files changed, 50 insertions(+), 74 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index ca0e8e7d04..911e42b0ee 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -198,17 +198,17 @@ def save_bounded_att_runtime_to_pages( page_kind, page_slot_ids: list[int], ) -> None: - """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + """把 prefill 基准 Conv 窗口和 canonical SSM 写入页。history 与 MTP 候选区不参与。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) - copy_bounded_att_tensors( - copies=[ - (self._accepted_conv(req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - ) - self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, :width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def restore_bounded_att_runtime_from_pages( self, @@ -344,24 +344,25 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayGDNState", ) -> None: - """复制已接受运行态到目标单槽，并清空目标 history。""" + """复制 prefill 基准 Conv/SSM 到目标单槽，并清空目标 history。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" assert ( type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config ), "DP copy requires the same algorithm and equivalent TP state configurations" width = self.config.conv_width - conv_copies = [] + conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): self.clear_history(dst_req) conv_copies.append( ( - remote_dp_bounded_state._accepted_conv(src_req), + remote_dp_bounded_state.conv_state[:, src_req, :, :width], self.conv_state[:, dst_req, :, :width], ) ) + ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + 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) @@ -384,11 +385,13 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + """写出 prefill 基准 Conv 窗口和 canonical SSM。源 history 不参与。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = self.get_accepted_state(req_idx) - fields = self._tensor_fields(accepted.conv, accepted.ssm) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -531,22 +534,6 @@ def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: to commit_ssm_prefix(indices, reqs, lengths) self.accept_updates(reqs, indices) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.ssm_state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = self.accepted_ssm(req) - dst.copy_(snap) - if flush: - self.ssm_state[:, req].copy_(snap) - self.clear_history(req) - # ======================================================================== # 七、私有辅助实现 # ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 874acdb070..1bed400125 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -198,17 +198,17 @@ def save_bounded_att_runtime_to_pages( page_kind, page_slot_ids: list[int], ) -> None: - """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + """把 prefill 基准 Conv 窗口和 canonical SSM 写入页。history 与 MTP 候选区不参与。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) - copy_bounded_att_tensors( - copies=[ - (self._accepted_conv(req_idx), pool.conv_state[slot_id]) - for req_idx, slot_id in zip(req_indices, page_slot_ids) - ] - ) - self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, :width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def restore_bounded_att_runtime_from_pages( self, @@ -344,24 +344,25 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayKDAState", ) -> None: - """复制已接受运行态到目标单槽，并清空目标 history。""" + """复制 prefill 基准 Conv/SSM 到目标单槽，并清空目标 history。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" assert ( type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config ), "DP copy requires the same algorithm and equivalent TP state configurations" width = self.config.conv_width - conv_copies = [] + conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): self.clear_history(dst_req) conv_copies.append( ( - remote_dp_bounded_state._accepted_conv(src_req), + remote_dp_bounded_state.conv_state[:, src_req, :, :width], self.conv_state[:, dst_req, :, :width], ) ) + ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) - remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + 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) @@ -384,11 +385,13 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + """写出 prefill 基准 Conv 窗口和 canonical SSM。源 history 不参与。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - accepted = self.get_accepted_state(req_idx) - fields = self._tensor_fields(accepted.conv, accepted.ssm) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -530,22 +533,6 @@ def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: to commit_ssm_prefix(indices, reqs, lengths) self.accept_updates(reqs, indices) - def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: - """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" - - if not req_indices: - return - target = targets[0] - if target.is_cuda and target.device == self.ssm_state.device: - self.export_accepted_state(req_indices, targets, flush=flush) - return - for req, dst in zip(req_indices, targets): - snap = self.accepted_ssm(req) - dst.copy_(snap) - if flush: - self.ssm_state[:, req].copy_(snap) - self.clear_history(req) - # ======================================================================== # 七、私有辅助实现 # ======================================================================== 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 index 0d0f1ec417..0eafbbbb20 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -643,7 +643,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg state.gates.uniform_(-0.1, -0.01) state.betas.uniform_(0.1, 0.9) active_before = state.ssm_state[:, 1].clone() - conv_expected = state.conv_state[:, 1, :, 1:4].clone() + conv_expected = state.conv_state[:, 1, :, :3].clone() block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") kind = "big" if save_big else "small" pool = state.big_pages_pool if save_big else state.small_pages_pool @@ -663,7 +663,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg 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.ssm_state[0], active_before.cpu(), rtol=0, atol=0) torch.testing.assert_close(pool.conv_state[0], conv_expected.cpu(), rtol=0, atol=0) torch.testing.assert_close(state.ssm_state[:, 1], active_before, rtol=0, atol=0) assert state.cursors[1].item() == count @@ -675,7 +675,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg state.mtp_state_index[2] = 2 restore() assert state.cursors[2].item() == state.mtp_state_index[2].item() == 0 - torch.testing.assert_close(state.ssm_state[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 2], active_before, rtol=0, atol=0) torch.testing.assert_close(state.conv_state[:, 2, :, :3], conv_expected, rtol=0, atol=0) state.reset_bounded_att_state([2]) assert not state.ssm_state[:, 2].any() diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index 47beee1eef..bda6fc61c1 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -109,10 +109,13 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( state.gates.fill_(-0.5) state.betas.zero_() state.cursors[1], state.cursors[2] = 1, 3 - if hasattr(state, "get_accepted_state"): - expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + width = state.config.conv_width + if mode == "replay": + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, 1].cpu().tolist(), + ] else: - width = state.config.conv_width expected = [ state.conv_state[:, 1, :, :width].cpu().tolist(), state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index a3a11d0b7b..e1c063dd50 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -53,7 +53,8 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( expected_conv = state.conv_state[:, src, ..., :3].clone() expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() else: - expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() + # prefill 页只保存基准窗口和 canonical SSM，不按 MTP 偏移取窗口，也不折回 history。 + expected_conv = state.conv_state[:, src, ..., :3].clone() expected_ssm = state.ssm_state[:, src].clone() if mode == "replay": state.raw_keys.zero_() @@ -61,7 +62,6 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( state.gates.fill_(-0.5) state.betas.zero_() state.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) @@ -130,12 +130,13 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty state.betas.zero_() state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) requests = [2, 0, 1] + width = state.config.conv_width if mode == "replay": - expected = [(item.conv.clone(), item.ssm.clone()) for item in map(state.get_accepted_state, requests)] + expected = [(state.conv_state[:, req, :, :width].clone(), state.ssm_state[:, req].clone()) for req in requests] else: expected = [ ( - state.conv_state[:, req, :, : state.config.conv_width].clone(), + state.conv_state[:, req, :, :width].clone(), state.ssm_state[:, req * (state.mtp_step + 1)].clone(), ) for req in requests diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index f583bcf7da..6da740353d 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -152,7 +152,6 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.gates.fill_(-0.5) linear.betas.zero_() linear.cursors[0] = 1 - ssm = linear.accepted_ssm(0).clone() linear.cursors[2] = 3 if mtp_step: linear.mtp_state_index[2] = mtp_step diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 7dbe235831..d230d58839 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -174,7 +174,6 @@ def test_pd_kv_and_runtime_state_roundtrip( backend.gates.fill_(-0.5) backend.betas.zero_() backend.cursors[src_req] = 1 - ssm[:, heads].copy_(mem.linear.accepted_ssm(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) From 8bfc2cc522b0d0b2d566d4b490e2c09d46b9678c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 08:28:59 +0000 Subject: [PATCH 35/54] refactor: drop the empty prefill fold from replay state Prefill writes the canonical SSM directly and leaves the cursor at zero. The snapshot helpers that only wrapped that empty prefix are gone. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 30 ------------------- .../bounded_att/kda/replay_kda_state.py | 29 ------------------ .../basemodel/attention/linear/replay_gdn.py | 1 - .../basemodel/attention/linear/replay_kda.py | 1 - .../test_hybrid_req_manager_pd_sharing.py | 10 +++---- unit_tests/common/test_linear_state.py | 20 ++++--------- 6 files changed, 9 insertions(+), 82 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 911e42b0ee..9ecd946c54 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -8,7 +8,6 @@ from functools import lru_cache from math import prod -from typing import NamedTuple import torch import triton @@ -31,11 +30,6 @@ logger = init_logger(__name__) -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - class ReplayGDNState(BoundedAttentionState): """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -239,10 +233,6 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: - """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) - # ======================================================================== # 三、CPU Cache 数据交换 # ======================================================================== @@ -432,17 +422,6 @@ def read_from_pd_page( # 六、Replay history # ======================================================================== - def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层把已接受前缀整段压进 SSM。""" - - layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: - return - if self.ssm_state.is_cuda: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" @@ -651,15 +630,6 @@ def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), } - def _accepted_conv(self, req_idx: int) -> torch.Tensor: - """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" - width = self.config.conv_width - conv = self.conv_state[:, req_idx] - if self.mtp_state_index is None: - return conv[:, :, :width] - offset = self.mtp_state_index[req_idx] - return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - def _materialize_accepted( self, reqs: torch.Tensor, diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 1bed400125..96522b7250 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -8,7 +8,6 @@ from functools import lru_cache from math import prod -from typing import NamedTuple import torch import triton @@ -31,11 +30,6 @@ logger = init_logger(__name__) -class _AcceptedTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - class ReplayKDAState(BoundedAttentionState): """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" @@ -239,10 +233,6 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to return self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) - def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: - """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" - return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) - # ======================================================================== # 三、CPU Cache 数据交换 # ======================================================================== @@ -431,16 +421,6 @@ def read_from_pd_page( # ======================================================================== # 六、Replay history # ======================================================================== - def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层把已接受前缀整段压进 SSM。""" - - layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: - return - if self.ssm_state.is_cuda: - self.merge_accepted_updates(reqs) - else: - self.clear_history(reqs) def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" @@ -650,15 +630,6 @@ def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), } - def _accepted_conv(self, req_idx: int) -> torch.Tensor: - """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" - width = self.config.conv_width - conv = self.conv_state[:, req_idx] - if self.mtp_state_index is None: - return conv[:, :, :width] - offset = self.mtp_state_index[req_idx] - return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) - def _materialize_accepted( self, reqs: torch.Tensor, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index d08a714d27..e88fc619e4 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -118,7 +118,6 @@ def prefill_att( conv_states = conv_states[:, :, : -backend.mtp_step] mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) - state.finish_prefill(layer_num, self.b_conv_buffer_idx) return core_attn_out, z def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 3090faa9a8..f5ccc66a9c 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -123,7 +123,6 @@ def prefill_att( lower_bound=backend.lower_bound, ) ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - state.finish_prefill(layer_num, self.b_conv_buffer_idx) return output diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index bda6fc61c1..ba349372ca 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -24,11 +24,9 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): 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) - if hasattr(state, "get_accepted_state"): - actual = state.get_accepted_state(2) - else: - width = state.config.conv_width - actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * (state.mtp_step + 1)]) + width = state.config.conv_width + ssm_index = 2 if mode == "replay" else 2 * (state.mtp_step + 1) + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, ssm_index]) for tensor, values in zip(actual, expected): torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) assert state.mtp_state_index[2] == 0 @@ -95,7 +93,7 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): @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 + runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization ): state, resources, _ = make_state(cls=cls, mode=mode) state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index e1c063dd50..02dfe2a6ca 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -85,11 +85,10 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( if mode == "replay": state.cursors[dst] = 3 restore() + actual_conv = state.conv_state[:, dst, ..., :3] if mode == "replay": - actual = state.get_accepted_state(dst) - actual_conv, actual_ssm = actual.conv, actual.ssm + actual_ssm = state.ssm_state[:, dst] else: - actual_conv = state.conv_state[:, dst, ..., :3] actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) @@ -116,7 +115,7 @@ def copy_on_cpu(*, copies): @pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("destination", ["cpu", "cuda"]) -def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dtype, destination): +def test_checkpoint_batch_export(runtime_start_args, cls, mode, dtype, destination): state, _, _ = make_state(cls, mode, 2, dtype) state.initialize_page_pools(3, device=destination) state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) @@ -144,14 +143,6 @@ def test_checkpoint_batch_export(runtime_start_args, monkeypatch, cls, mode, dty before = state.ssm_state.clone() history = getattr(state, "runtime", state) cursors = history.cursors.clone() if mode == "replay" else None - - def no_snapshot(*args): - raise AssertionError("batch export must write directly to the destination") - - if mode == "replay": - monkeypatch.setattr(state, "get_accepted_state", no_snapshot) - if history is not state: - monkeypatch.setattr(history, "get_accepted_state", no_snapshot) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): @@ -169,11 +160,10 @@ def no_snapshot(*args): target.mtp_state_index.fill_(2) target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) for req, (conv, ssm) in enumerate(expected): + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] if mode == "replay": - actual = target.get_accepted_state(req) - actual_conv, actual_ssm = actual.conv, actual.ssm + actual_ssm = target.ssm_state[:, req] else: - actual_conv = target.conv_state[:, req, :, : target.config.conv_width] actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) From 28db98707d3733b8b1564fc98caff2908ff75364 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 10:49:30 +0000 Subject: [PATCH 36/54] refactor: advance the replay cursor after the whole window is accepted MTP verify trims the rejected tail from that cursor, so decode no longer records a per-step write base. Co-authored-by: Cursor --- .../bounded_att/gdn/gdn_state.py | 1 + .../bounded_att/gdn/replay_gdn_state.py | 201 +++++------------- .../attention_state/bounded_att/interface.py | 6 +- .../bounded_att/kda/kda_state.py | 1 + .../bounded_att/kda/replay_kda_state.py | 201 +++++------------- .../sliding_window/fp8_e8m0_swa_state.py | 7 +- .../window_compressor_state.py | 7 +- lightllm/common/attention_state/manager.py | 12 +- .../basemodel/attention/linear/replay_gdn.py | 3 +- .../attention/linear/replay_gdn_decode.py | 1 - .../basemodel/attention/linear/replay_kda.py | 4 +- .../attention/linear/replay_kda_decode.py | 1 - .../triton_kernel/linear_att/replayssm.py | 4 +- .../linear_att/replayssm_cursor.py | 43 +++- .../linear_att/replayssm_decode.py | 11 +- .../linear_att/replayssm_fold.py | 45 ++++ .../linear_att/ssm_acceptance.py | 58 +++-- .../triton_kernel/linear_att/ssm_autotune.py | 13 +- .../model_infer/mtp_speculative/utils.py | 22 +- test/att_resources/validate_model_cache.py | 2 + test/benchmark/kernels/benchmark_kda_ssm.py | 8 +- .../att_resources/test_backend_startup.py | 2 + .../att_resources/test_bounded_state.py | 6 +- .../att_resources/test_mtp_state_commit.py | 5 + .../linear_att/replay_history.py | 50 ++++- .../linear_att/test_acceptance_fusion.py | 65 +++--- .../linear_att/test_replayssm.py | 74 +++---- .../linear_att/test_replayssm_execution.py | 42 ++-- .../common/test_linear_state_contract.py | 42 +++- unit_tests/models/glm5_next/test_mtp.py | 12 +- 30 files changed, 492 insertions(+), 457 deletions(-) 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..b001d74ae3 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -178,6 +178,7 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 9ecd946c54..9506fc6d4e 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -10,12 +10,9 @@ from math import prod import torch -import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( - _advance_cursor, - _fold_prefix, -) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor as advance_replay_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -104,10 +101,8 @@ def __init__( assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 self.hold_req_id = slots - 1 - # cursor 是已接受、尚未折回的长度，各层只读。 - # write_bases 是本步 replay 的写入起点：层内折回后为 0，否则等于进入 kernel 时的 cursor。 + # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) - self.write_bases = torch.zeros(slots, dtype=torch.int32, device=device) # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # GDN 的 g 是每个头一个标量。 @@ -181,7 +176,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: for req_idx in req_indices: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() - self.clear_history(req_idx) + self.cursors[req_idx].zero_() if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) @@ -211,14 +206,13 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + """从页恢复基准 Conv/SSM。prefill 时 cursor 仍是 0。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) width = self.config.conv_width conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - self.clear_history(req_idx) conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) @@ -227,11 +221,16 @@ def restore_bounded_att_runtime_from_pages( 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: - """提交已接受长度。前缀放不下时，折回发生在下一次 replay kernel 里。""" + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: + """提交已接受长度，并剪掉整窗全接受后未通过验证的尾部。""" if self.mtp_state_index is None: return - self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps, req_verify_width) # ======================================================================== # 三、CPU Cache 数据交换 @@ -334,7 +333,7 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayGDNState", ) -> None: - """复制 prefill 基准 Conv/SSM 到目标单槽，并清空目标 history。""" + """复制 prefill 基准 Conv/SSM 到目标单槽。prefill 时目标 cursor 仍是 0。""" 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 ( @@ -343,7 +342,6 @@ def copy_runtime_state_from_dp( width = self.config.conv_width conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.clear_history(dst_req) conv_copies.append( ( remote_dp_bounded_state.conv_state[:, src_req, :, :width], @@ -399,10 +397,9 @@ def read_from_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """读入基准 Conv/SSM，并清空该请求的 history。""" + """读入基准 Conv/SSM。prefill 时 cursor 仍是 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" - self.clear_history(req_idx) fields = self._tensor_fields( self.conv_state[:, req_idx, :, : self.config.conv_width], self.ssm_state[:, req_idx], @@ -420,98 +417,65 @@ def read_from_pd_page( # ======================================================================== # 六、Replay history + # + # cursor 是已接受、尚未折回的长度，各层只读。 + # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 # ======================================================================== - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" + def finish_decode_token( + self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None + ) -> None: + """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 - self.accept_updates(reqs, reqs, width=1) + if cu_seqlens is None: + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + self.advance_cursor(reqs, cu_seqlens) - def clear_history(self, req: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + """按 cu_seqlens 给出的整窗全接受推进 cursor，不修改 SSM。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - self.write_bases[req].zero_() - else: - self.cursors[req] = 0 - self.write_bases[req] = 0 + advance_replay_cursor(self, reqs, cu_seqlens) - def accept_updates( - self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + def commit_accepted( + self, + mtp_state_index: torch.Tensor, + reqs: torch.Tensor, + accepted_lengths: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: - """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" - - _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - self.write_bases, - reqs, - accepted_state_indices, - reqs.numel(), - self.hold_req_id, - self.verify_width if width is None else width, - 256, - ) + """MTP verify 之后写入索引，并从全接受后的 cursor 上剪掉未接受尾部。 - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 + 配置的 verify_width 为 1 时没有尾部，只把这一个 token 记入 cursor。 + """ - if reqs.numel() == 0: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False) - self.cursors.index_fill_(0, reqs.long(), 0) - - def accepted_ssm(self, req_idx: int) -> torch.Tensor: - """返回不修改 canonical SSM 和 history 的单请求快照。""" + if self.verify_width > 1: + commit_ssm_prefix(mtp_state_index, reqs, accepted_lengths, req_verify_width, self.cursors) + else: + mtp_state_index[reqs] = (accepted_lengths - 1).to(mtp_state_index.dtype) + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + self.advance_cursor(reqs, cu_seqlens) - layers, _, hv, k, v = self.ssm_state.shape - output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) - reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) - self._materialize_accepted(reqs, output, snapshot=True) - return output + def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 - def export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页。""" + snapshot 为空时写回 canonical SSM，并清掉这些请求的 cursor。 + 传入 snapshot 时只写入这块独立缓冲，不改 canonical SSM，也不动 cursor。 + run_config 为空时由 fold_prefix 使用它自己的默认配置。 + """ - if not req_indices: + if reqs.numel() == 0: return - shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) - device = self.ssm_state.device - self._materialize_accepted( - torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), - torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( - device, non_blocking=True - ), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - if flush: - slots = torch.tensor(req_indices, dtype=torch.int64, device=device) - self.cursors.index_fill_(0, slots, 0) - self.write_bases.index_fill_(0, slots, 0) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并把 cursor 设为这次 replay 的写起点加接受长度。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors, self.write_bases) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) + if snapshot is None: + fold_prefix(self, reqs, self.ssm_state, False, run_config) + self.cursors.index_fill_(0, reqs.long(), 0) + return + fold_prefix(self, reqs, snapshot, True, run_config) # ======================================================================== # 七、私有辅助实现 @@ -629,52 +593,3 @@ def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), } - - def _materialize_accepted( - self, - reqs: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - layers, slots, hv, k, v = self.ssm_state.shape - layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() - # 折回不走 replay 的 autotune。没固定配置时用这份默认值。 - config = self.run_config or {} - bv = config.get("BV", 32) - num_warps = config.get("num_warps", 4) - num_stages = config.get("num_stages", 3) - _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( - self.ssm_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, - False, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - STRIDE_LAYER=layer_stride, - STRIDE_SLOT=slot_stride, - STRIDE_HEAD=head_stride, - STRIDE_K=key_stride, - STRIDE_V=value_stride, - num_warps=num_warps, - num_stages=num_stages, - ) diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index fe1fceb3bd..11d7425cd9 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -102,8 +102,12 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。 + + req_verify_width 与请求逐行对齐，是每个请求这次送去验证的 token 数。Replay 用它从全接受后的 cursor 上剪掉未接受尾部。 + """ ... diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index e5480a3ab0..9e386dbe2d 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -178,6 +178,7 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 96522b7250..daa297dbb9 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -10,12 +10,9 @@ from math import prod import torch -import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( - _advance_cursor, - _fold_prefix, -) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor as advance_replay_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -104,10 +101,8 @@ def __init__( assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 self.hold_req_id = slots - 1 - # cursor 是已接受、尚未折回的长度，各层只读。 - # write_bases 是本步 replay 的写入起点：层内折回后为 0，否则等于进入 kernel 时的 cursor。 + # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) - self.write_bases = torch.zeros(slots, dtype=torch.int32, device=device) # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # KDA 的 g 按 key 通道存。 @@ -181,7 +176,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: for req_idx in req_indices: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() - self.clear_history(req_idx) + self.cursors[req_idx].zero_() if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) @@ -211,14 +206,13 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + """从页恢复基准 Conv/SSM。prefill 时 cursor 仍是 0。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) width = self.config.conv_width conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - self.clear_history(req_idx) conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) @@ -227,11 +221,16 @@ def restore_bounded_att_runtime_from_pages( 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: - """提交已接受长度。前缀放不下时，折回发生在下一次 replay kernel 里。""" + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: + """提交已接受长度，并剪掉整窗全接受后未通过验证的尾部。""" if self.mtp_state_index is None: return - self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps, req_verify_width) # ======================================================================== # 三、CPU Cache 数据交换 @@ -334,7 +333,7 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayKDAState", ) -> None: - """复制 prefill 基准 Conv/SSM 到目标单槽，并清空目标 history。""" + """复制 prefill 基准 Conv/SSM 到目标单槽。prefill 时目标 cursor 仍是 0。""" 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 ( @@ -343,7 +342,6 @@ def copy_runtime_state_from_dp( width = self.config.conv_width conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.clear_history(dst_req) conv_copies.append( ( remote_dp_bounded_state.conv_state[:, src_req, :, :width], @@ -399,10 +397,9 @@ def read_from_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """读入基准 Conv/SSM，并清空该请求的 history。""" + """读入基准 Conv/SSM。prefill 时 cursor 仍是 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" - self.clear_history(req_idx) fields = self._tensor_fields( self.conv_state[:, req_idx, :, : self.config.conv_width], self.ssm_state[:, req_idx], @@ -420,98 +417,65 @@ def read_from_pd_page( # ======================================================================== # 六、Replay history + # + # cursor 是已接受、尚未折回的长度，各层只读。 + # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 # ======================================================================== - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: - """最后一层记下这 1 个 token。前缀放不下时，折回发生在下一次 replay kernel 里。""" + def finish_decode_token( + self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None + ) -> None: + """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 - self.accept_updates(reqs, reqs, width=1) + if cu_seqlens is None: + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + self.advance_cursor(reqs, cu_seqlens) - def clear_history(self, req: int | torch.Tensor) -> None: - """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + """按 cu_seqlens 给出的整窗全接受推进 cursor，不修改 SSM。""" - if type(req) is int: - # 标量赋值会把 CPU tensor 复制到 CUDA slot。 - self.cursors[req].zero_() - self.write_bases[req].zero_() - else: - self.cursors[req] = 0 - self.write_bases[req] = 0 + advance_replay_cursor(self, reqs, cu_seqlens) - def accept_updates( - self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + def commit_accepted( + self, + mtp_state_index: torch.Tensor, + reqs: torch.Tensor, + accepted_lengths: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: - """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" - - _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - self.write_bases, - reqs, - accepted_state_indices, - reqs.numel(), - self.hold_req_id, - self.verify_width if width is None else width, - 256, - ) + """MTP verify 之后写入索引，并从全接受后的 cursor 上剪掉未接受尾部。 - def merge_accepted_updates(self, reqs: torch.Tensor) -> None: - """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 + 配置的 verify_width 为 1 时没有尾部，只把这一个 token 记入 cursor。 + """ - if reqs.numel() == 0: - return - self._materialize_accepted(reqs, self.ssm_state, snapshot=False) - self.cursors.index_fill_(0, reqs.long(), 0) - - def accepted_ssm(self, req_idx: int) -> torch.Tensor: - """返回不修改 canonical SSM 和 history 的单请求快照。""" + if self.verify_width > 1: + commit_ssm_prefix(mtp_state_index, reqs, accepted_lengths, req_verify_width, self.cursors) + else: + mtp_state_index[reqs] = (accepted_lengths - 1).to(mtp_state_index.dtype) + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + self.advance_cursor(reqs, cu_seqlens) - layers, _, hv, k, v = self.ssm_state.shape - output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) - reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) - self._materialize_accepted(reqs, output, snapshot=True) - return output + def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 - def export_accepted_state( - self, - req_indices: list[int], - targets: list[torch.Tensor], - flush: bool = False, - ) -> None: - """将多个 accepted 快照直接写入调用方目标页。""" + snapshot 为空时写回 canonical SSM，并清掉这些请求的 cursor。 + 传入 snapshot 时只写入这块独立缓冲，不改 canonical SSM，也不动 cursor。 + run_config 为空时由 fold_prefix 使用它自己的默认配置。 + """ - if not req_indices: + if reqs.numel() == 0: return - shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) - strides = targets[0].stride() - assert len(req_indices) == len(targets) - assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) - device = self.ssm_state.device - self._materialize_accepted( - torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), - torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( - device, non_blocking=True - ), - snapshot=not flush, - flush_export=flush, - output_strides=strides, - ) - if flush: - slots = torch.tensor(req_indices, dtype=torch.int64, device=device) - self.cursors.index_fill_(0, slots, 0) - self.write_bases.index_fill_(0, slots, 0) - - def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: - """把接受长度写入 MTP 槽位索引，并把 cursor 设为这次 replay 的写起点加接受长度。""" - - if self.verify_width > 1: - commit_ssm_prefix(indices, reqs, lengths, self.cursors, self.write_bases) - else: - commit_ssm_prefix(indices, reqs, lengths) - self.accept_updates(reqs, indices) + if snapshot is None: + fold_prefix(self, reqs, self.ssm_state, False, run_config) + self.cursors.index_fill_(0, reqs.long(), 0) + return + fold_prefix(self, reqs, snapshot, True, run_config) # ======================================================================== # 七、私有辅助实现 @@ -629,52 +593,3 @@ def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), } - - def _materialize_accepted( - self, - reqs: torch.Tensor, - output: torch.Tensor, - snapshot: bool, - output_strides=None, - flush_export: bool = False, - ) -> None: - """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" - - assert not (snapshot and flush_export), "flush export must update the canonical state" - layers, slots, hv, k, v = self.ssm_state.shape - layer_stride, slot_stride, head_stride, key_stride, value_stride = self.ssm_state.stride() - # 折回不走 replay 的 autotune。没固定配置时用这份默认值。 - config = self.run_config or {} - bv = config.get("BV", 32) - num_warps = config.get("num_warps", 4) - num_stages = config.get("num_stages", 3) - _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( - self.ssm_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, - True, - FLUSH_EXPORT=flush_export, - OUTPUT_STRIDES=output_strides, - STRIDE_LAYER=layer_stride, - STRIDE_SLOT=slot_stride, - STRIDE_HEAD=head_stride, - STRIDE_K=key_stride, - STRIDE_V=value_stride, - num_warps=num_warps, - num_stages=num_stages, - ) diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index a1d72102ac..5afdfab160 100644 --- a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py +++ b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py @@ -179,7 +179,12 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: """ring 按绝对 token 位置寻址，无需提交额外游标。""" pass diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index cea0a49af2..5d9a0ffe12 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -210,7 +210,12 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..c0a85c367e 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -396,13 +396,19 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, role: str, ) -> None: - """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + """传入 CUDA 请求索引、int32 接受步数和每个请求的验证宽度；无参与 State 时为空操作。 + + req_verify_width 与请求逐行对齐，是每个请求这次送去验证的 token 数。Replay 用它剪掉未接受尾部。 + """ + assert ( + req_indices.is_cuda and accepted_steps.is_cuda and req_verify_width.is_cuda + ), "commit inputs must be CUDA tensors" assert role in MODEL_ROLES, f"unknown model role {role}" for state in self.get_bounded_attention_states(role): - state.commit_bounded_att_state(req_indices, accepted_steps) + state.commit_bounded_att_state(req_indices, accepted_steps, req_verify_width) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index e88fc619e4..eb54adcfaa 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -190,7 +190,8 @@ def decode_att( core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) else: core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) - state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) return core, z def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 36b589563e..8bfbcf07ea 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -61,7 +61,6 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, - history.write_bases, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index f5ccc66a9c..a03f167563 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -195,8 +195,8 @@ def decode_att( self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, ) - if backend.mtp_step == 0: - state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) return output def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index de99c5fba4..e9f8c2072f 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -61,7 +61,6 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, - history.write_bases, cu_seqlens, out, q.stride(axis), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index c0a598143e..7b57371163 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """ReplaySSM kernel exports, in lifecycle order. -decode writes history for the current window. -cursor records how many of those tokens were accepted. +decode writes history for the current window and does not move the cursor. +advance_cursor moves the cursor after the last layer; verify trims the rejected tail. fold applies that accepted prefix onto the canonical SSM. """ diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py index 5eee325083..ea3121b64b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -"""Advance the accepted-prefix cursor. This kernel does not fold state.""" +"""Update the accepted-prefix cursor after replay has returned. + +Replay itself does not touch the cursor: every layer must still see the length +from the start of this step. After the last layer returns, advance the cursor +over the replayed window. MTP verify then trims the rejected tail. +""" import triton import triton.language as tl @@ -8,18 +13,40 @@ @triton.jit(do_not_specialize=["N"]) def _advance_cursor( cursors, - write_bases, reqs, - accepted_state_indices, + cu, N, HOLD: tl.constexpr, - WIDTH: tl.constexpr, + CAPACITY: tl.constexpr, BLOCK: tl.constexpr, ): - # 写起点由 replay kernel 记下。半区已经折回时起点是 0，否则是原 cursor。 + # 先按整窗全接受。放不下时新 token 占 [0, window)，cursor 记为 window；否则加上 window。 i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) r = tl.load(reqs + i, i < N, HOLD) valid = (i < N) & (r != HOLD) - base = tl.load(write_bases + r, valid, 0) - count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 - tl.store(cursors + r, base + count, valid) + cursor = tl.load(cursors + r, valid, 0) + window = tl.load(cu + i + 1, valid, 0) - tl.load(cu + i, valid, 0) + updated = tl.where(cursor + window > CAPACITY, window, cursor + window) + tl.store(cursors + r, updated, valid) + + +def advance_cursor(history, reqs, cu_seqlens) -> None: + """最后一层返回后，把 cursor 推进到整窗都被接受的位置。 + + cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。窗口长度为 0 时 cursor 保持不变。HOLD 请求跳过。 + """ + + if reqs.numel() == 0: + return + assert reqs.is_cuda and cu_seqlens.is_cuda, "advance_cursor requires CUDA inputs" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + history.cursors, + reqs, + cu_seqlens, + reqs.numel(), + history.hold_req_id, + history.capacity, + 256, + num_warps=4, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index df76b0ab20..fae106be22 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -46,7 +46,6 @@ def _decode_replay( Betas, Reqs, Cursors, - WriteBase, Cu, Out, SQ: tl.constexpr, @@ -74,7 +73,7 @@ def _decode_replay( """执行一个请求窗口的 replay decode，并写入 history。 前缀放不下本步 token 时，先在本 kernel 里把已有前缀从原始 k/v 折进 checkpoint， - 舍入成 SSM 的 dtype，再从这份状态重放本步。cursor 本 kernel 不修改，写起点记到 WriteBase。 + 舍入成 SSM 的 dtype，再从这份状态重放本步。cursor 本 kernel 不修改，留给最后一层返回后更新。 """ iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) req = tl.load(Reqs + seq).to(tl.int64) @@ -87,12 +86,10 @@ def _decode_replay( for t in range(start, end): tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) return - n = tl.load(Cursors + req) if start == end: - # 本步没有 token。写起点保持原 cursor，接受 0 个 token 时长度不变。 - if iv == 0 and hv == 0: - tl.store(WriteBase + req, n) + # 本步没有 token。cursor 保持原值，接受 0 个 token 时长度不变。 return + n = tl.load(Cursors + req) kk = tl.arange(0, BK) ll = tl.arange(0, L) @@ -106,8 +103,6 @@ def _decode_replay( state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) n = 0 - if iv == 0 and hv == 0: - tl.store(WriteBase + req, n) if PRECOMPUTE_PROJECTIONS: # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index e8742846aa..b0a9e48dda 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -86,3 +86,48 @@ def _fold_prefix( if FLUSH_EXPORT: tl.store(sp, state, mask) # cursor 不在这里清。和折回并行的其他线程块还要读原来的长度，清零放到 kernel 结束之后。 + + +def fold_prefix(history, reqs, output, snapshot: bool, run_config=None) -> None: + """把 history 上 cursor 覆盖的已接受前缀折进 output。 + + snapshot 为真时 output 是独立缓冲，kernel 不写回 canonical SSM。 + 没传入 run_config 时用 BV=32、4 个 warp、3 个 stage。 + """ + + if run_config is None: + run_config = {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = run_config["BV"] + num_warps = run_config["num_warps"] + num_stages = run_config["num_stages"] + state = history.ssm_state + layers, slots, hv, k, v = state.shape + layer_stride, slot_stride, head_stride, key_stride, value_stride = state.stride() + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + state, + output, + history.raw_keys, + history.raw_values, + history.gates, + history.betas, + history.cursors, + reqs, + slots, + hv, + history.num_key_heads, + k, + v, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + history.kda, + STRIDE_LAYER=layer_stride, + STRIDE_SLOT=slot_stride, + STRIDE_HEAD=head_stride, + STRIDE_K=key_stride, + STRIDE_V=value_stride, + num_warps=num_warps, + num_stages=num_stages, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py index 86ad1c8025..bf4a61e95a 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -6,62 +6,54 @@ def commit_ssm_prefix( - indices: torch.Tensor, + mtp_state_index: torch.Tensor, reqs: torch.Tensor, - lengths: torch.Tensor, - replay_cursors: torch.Tensor | None = None, - write_bases: torch.Tensor | None = None, + accepted_lengths: torch.Tensor, + req_verify_width: torch.Tensor, + replay_cursors: torch.Tensor, ) -> None: - """把每个请求的 accepted length 写入 MTP 索引。 + """写入 MTP 索引，并从已按验证宽度全接受推进的 cursor 上剪掉未接受尾部。 - 传入 write_bases 时，cursor 设为这次 replay 的写起点加接受长度。层内已经折回时写起点是 0。 + accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 """ - if not reqs.is_cuda: - indices[reqs] = (lengths - 1).to(indices.dtype) - if replay_cursors is not None: - valid = reqs != replay_cursors.numel() - 1 - accepted = lengths[valid].to(replay_cursors.dtype) - if write_bases is not None: - replay_cursors[reqs[valid]] = write_bases[reqs[valid]] + accepted - else: - replay_cursors[reqs[valid]] += accepted - return + + assert reqs.is_cuda, "commit_ssm_prefix requires CUDA requests" if reqs.numel(): _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( - indices, + mtp_state_index, reqs, - lengths, + accepted_lengths, replay_cursors, - write_bases if write_bases is not None else (replay_cursors if replay_cursors is not None else indices), + req_verify_width, reqs.numel(), - indices.numel() - 1, + mtp_state_index.numel() - 1, reqs.stride(0), - lengths.stride(0), + accepted_lengths.stride(0), + req_verify_width.stride(0), 256, - WRITE_BASE=write_bases is not None, num_warps=4, ) @triton.jit(do_not_specialize=["N"]) def _commit_prefix( - Indices, + MtpIndex, Reqs, - Lengths, + Accepted, Cursors, - WriteBase, + ReqVerifyWidth, N, HOLD: tl.constexpr, REQ_STRIDE: tl.constexpr, - LENGTH_STRIDE: tl.constexpr, + ACCEPTED_STRIDE: tl.constexpr, + VERIFY_STRIDE: tl.constexpr, BLOCK: tl.constexpr, - WRITE_BASE: tl.constexpr, ): row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) - length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) - tl.store(Indices + req, length - 1, row < N) - if Cursors is not None: - valid = (row < N) & (req != HOLD) - base = tl.load(WriteBase + req, valid, 0) if WRITE_BASE else tl.load(Cursors + req, valid, 0) - tl.store(Cursors + req, base + length, valid) + accepted = tl.load(Accepted + row * ACCEPTED_STRIDE, row < N, 0) + tl.store(MtpIndex + req, accepted - 1, row < N) + valid = (row < N) & (req != HOLD) + verify_tokens = tl.load(ReqVerifyWidth + row * VERIFY_STRIDE, valid, 0) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor - (verify_tokens - accepted), valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 2864f4f318..59b1ff5583 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -109,8 +109,6 @@ def rebuild_inputs( ) ) 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))) return (), dict( ssm_backend=scratch, mode=mode, @@ -122,7 +120,7 @@ def rebuild_inputs( a_log=inputs[5], bias=inputs[6], cu_seqlens=cu if cu_seqlens is not None else None, - workload=(reqs, accepted), + workload=(reqs,), ) @@ -158,14 +156,17 @@ def select_config( if workload is None: return run_config ssm_backend.run_config = run_config - reqs, accepted = workload + (reqs,) = workload ssm_backend.cursors.zero_() - # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。 + # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。调优窗口全部接受。 tokens = q.shape[1 if cu_seqlens is not None else 0] rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) for _ in range(rounds): replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) - ssm_backend.accept_updates(reqs, accepted) + adopt_cu = cu_seqlens + if adopt_cu is None: + adopt_cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + ssm_backend.advance_cursor(reqs, adopt_cu) return run_config diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index a32cedd17d..19d4656f14 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -32,12 +32,32 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) + # 最后一层已经按这个窗口全接受推进了 cursor。这里剪掉没通过验证的尾部。 backend.model.att_state_manager.commit_bounded_att_state( - b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" + b_req_idx[b_req_mtp_start_loc], + accept_lengths, + role="main", + req_verify_width=_req_verify_width(b_req_mtp_start_loc, b_req_idx.shape[0]), ) return accept_lengths, accepted_index +def _req_verify_width(starts: torch.Tensor, num_verify_tokens: int) -> torch.Tensor: + """每个请求这次送去验证的 token 数，与 b_req_mtp_start_loc 逐行对齐。 + + num_verify_tokens 是验证 batch 的 token 数，用作最后一个请求的结束位置。 + TODO: 并进 MTP verify 的同一个算子，直接用 b_req_mtp_start_loc 裁剪 cursor，省掉这次单独的宽度计算。 + """ + + req_verify_width = torch.empty_like(starts) + if starts.numel() == 0: + return req_verify_width + if starts.numel() > 1: + req_verify_width[:-1].copy_(starts[1:] - starts[:-1]) + req_verify_width[-1] = num_verify_tokens - starts[-1] + return req_verify_width + + def scatter_mtp_next_tokens( backend: ModeBackend, proposal: SpecProposal, # proposal.token_ids: [req_num, draft_step] diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 4bf1514704..569d4f99a5 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -236,10 +236,12 @@ def decode(request: torch.Tensor): if draft_models and hybrid: cuda_req_indices = request.cuda(non_blocking=True) accepted_steps = torch.full_like(cuda_req_indices, accepted_width, dtype=torch.int32) + req_verify_width = torch.full_like(cuda_req_indices, step_width, dtype=torch.int32) for item in (model, *draft_models): resources.commit_bounded_att_state( req_indices=cuda_req_indices, accepted_steps=accepted_steps, + req_verify_width=req_verify_width, role=item.att_state_role, ) return logits diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index c50adef255..706807de5a 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -16,6 +16,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, attach_replay_history as ReplaySSMExecutor, ) @@ -73,10 +74,13 @@ def cycle(): else: replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, cu) if ssm_backend is not None: - ssm_backend.accept_updates(reqs, accepted) + if cu is None: + accept_replay_window(ssm_backend, reqs) + else: + accept_replay_window(ssm_backend, reqs, accepted, cu) if isinstance(ssm_backend, ReplaySSMExecutor): ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) - ssm_backend.merge_accepted_updates(reqs) + ssm_backend.fold_accepted_prefix(reqs) else: snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e3a09adb8a..457fdfa1c3 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -298,6 +298,7 @@ def start_runtime(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), role="main", ) assert all(not state.mtp_state_index.any() for state in draft_states) @@ -308,6 +309,7 @@ def start_runtime(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([i + 2], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), role=f"draft{i}", ) assert state.mtp_state_index[0] == i + 1 diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index ffd59db4aa..e2c08b1461 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -254,7 +254,11 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): assert (state.small_pages_pool.conv_state[0] == 7).all() assert (state.small_pages_pool.ssm_state[0] == 11).all() with pytest.raises(AssertionError, match="CUDA tensors"): - state.commit_bounded_att_state(torch.tensor([0]), torch.tensor([1], dtype=torch.int32)) + state.commit_bounded_att_state( + torch.tensor([0]), + torch.tensor([1], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + ) @pytest.mark.parametrize("cls", [KDAState, GDNState]) diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index a50ff4fc71..ea114d2ba9 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -36,6 +36,7 @@ def test_commit_selects_state_for_one_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="main", ) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.main_gdn_conv_fp32_ssm_fp32_state): @@ -46,6 +47,7 @@ def test_commit_selects_state_for_one_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([1], dtype=torch.int32, device="cuda"), + torch.tensor([1], dtype=torch.int32, device="cuda"), role="draft0", ) assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 1] @@ -60,17 +62,20 @@ def test_commit_without_recurrent_state_and_invalid_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="main", ) manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="draft0", ) with pytest.raises(AssertionError, match="unknown model role"): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="target", ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 916291e26b..ede0beae91 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -1,6 +1,8 @@ """测试用的 replay history。正式路径走 ReplayGDNState / ReplayKDAState 的构造函数。""" import torch +import triton +import triton.language as tl from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState @@ -51,10 +53,9 @@ def attach_replay_history( layers, slots, hv, k, v = ssm.shape owner.num_key_heads = hv if num_key_heads is None else num_key_heads assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 - # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度，write_bases 是本步写入起点。 + # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度。 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.write_bases = torch.zeros(slots, dtype=torch.int32, device=ssm.device) # 重建用 k/d：bf16 激活用 fp16。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) @@ -66,3 +67,48 @@ def attach_replay_history( owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) return owner + + +@triton.jit(do_not_specialize=["N"]) +def _trim_rejected_index( + cursors, + reqs, + cu, + accepted_index, + N, + HOLD: tl.constexpr, + BLOCK: tl.constexpr, +): + # accepted_index 按请求槽存放，接受个数是 index + 1。窗口长度与 reqs 逐行对齐。 + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + window = tl.load(cu + i + 1, valid, 0) - tl.load(cu + i, valid, 0) + accepted = tl.load(accepted_index + r, valid, 0) + 1 + cursor = tl.load(cursors + r, valid, 0) + tl.store(cursors + r, cursor - (window - accepted), valid) + + +def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens=None) -> None: + """测试里一次结算 cursor：先按整窗全接受，再按槽位 index 剪掉未接受尾部。 + + 没传 cu_seqlens 时每个请求的窗口长度是 1。正式路径用 advance_cursor 和 commit_accepted。 + """ + + if reqs.numel() == 0: + return + if cu_seqlens is None: + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + history.advance_cursor(reqs, cu_seqlens) + if cu_seqlens is None or accepted_state_indices is None: + return + _trim_rejected_index[(triton.cdiv(reqs.numel(), 256),)]( + history.cursors, + reqs, + cu_seqlens, + accepted_state_indices, + reqs.numel(), + history.hold_req_id, + 256, + num_warps=4, + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index a826fb786b..d384fadd92 100644 --- 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 @@ -1,20 +1,22 @@ import pytest import torch -import triton from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, attach_replay_history as ReplaySSMExecutor, ) def ssm_snapshot(backend, req): - return backend.accepted_ssm(req) + state = backend.ssm_state + output = torch.empty((state.shape[0], *state.shape[2:]), dtype=state.dtype, device=state.device) + backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device=state.device), snapshot=output) + return output from types import SimpleNamespace from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer @@ -33,9 +35,9 @@ def make(n, width): args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] ssm_backend = ReplayGDNState.__new__(ReplayGDNState) ssm_backend.cursors = tensor([7] * (hold + 1)) - ssm_backend.write_bases = ssm_backend.cursors.clone() ssm_backend.hold_req_id = hold - ssm_backend.capacity = max(4, triton.next_power_of_2(width)) + # 这个用例不触发折回，容量要盖住多步累加后的 cursor。 + ssm_backend.capacity = 4096 ssm_backend.verify_width = width manager = SimpleNamespace() manager.mtp_state_index = tensor([-9] * (hold + 1)) @@ -43,15 +45,22 @@ def make(n, width): return manager, args, (ids, lengths, starts) +def _replay_cu(args): + """args[0] 是每个请求的起点，最后一个请求接到本批 token 末尾。""" + + starts = args[0] + return torch.cat((starts, starts.new_full((1,), args[3].numel()))) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("replay", [False, True]) -def test_request_counts_share_compiled_kernel(monkeypatch, replay): +def test_request_counts_share_compiled_kernel(monkeypatch): from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix reqs = torch.arange(512, dtype=torch.int32, device="cuda") lengths = reqs % 4 + 1 + window = lengths + 4 indices = torch.zeros(513, dtype=torch.int32, device="cuda") - cursors = torch.zeros_like(indices) if replay else None + cursors = torch.zeros_like(indices) hashes = set() run = _commit_prefix.run @@ -63,20 +72,23 @@ def record(*args, **kwargs): monkeypatch.setattr(_commit_prefix, "run", record) for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): indices.fill_(-1) - if replay: - cursors.zero_() - commit_ssm_prefix(indices, reqs[:count], lengths[:count], cursors) + cursors.zero_() + cursors[:count] = window[:count] + commit_ssm_prefix(indices, reqs[:count], lengths[:count], window[:count], cursors) assert torch.equal(indices[:count], lengths[:count] - 1) assert (indices[count:] == -1).all() - if replay: - assert torch.equal(cursors[:count], lengths[:count]) - assert not cursors[count:].any() + assert torch.equal(cursors[:count], lengths[:count]) + assert not cursors[count:].any() assert len(hashes) == 1 def accept_unfused(m, args): update(m.mtp_state_index, *args) - m._ssm_executor.accept_updates(args[1][args[0].long()], m.mtp_state_index) + reqs = args[1][args[0].long()] + if m._ssm_executor.verify_width > 1: + accept_replay_window(m._ssm_executor, reqs, m.mtp_state_index, _replay_cu(args)) + else: + accept_replay_window(m._ssm_executor, reqs) def accept_fused(m, args): @@ -84,7 +96,13 @@ def accept_fused(m, args): ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) lengths = cumulative[ends.long()] - cumulative[args[0].long()] - m._ssm_executor.commit_accepted(m.mtp_state_index, reqs, lengths) + executor = m._ssm_executor + if getattr(executor, "verify_width", 1) > 1: + cu = _replay_cu(args) + executor.advance_cursor(reqs, cu) + executor.commit_accepted(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1]) + else: + executor.commit_accepted(m.mtp_state_index, reqs, lengths, lengths) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -107,8 +125,6 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): ref.mtp_state_index.fill_(-9) m._ssm_executor.cursors.fill_(7) ref._ssm_executor.cursors.fill_(7) - m._ssm_executor.write_bases.fill_(7) - ref._ssm_executor.write_bases.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): @@ -118,9 +134,6 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): if r != m._ssm_executor.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) - # replay 没有折回时，写起点就是当前 cursor。提交把它加上本步接受数。 - m._ssm_executor.write_bases.copy_(m._ssm_executor.cursors) - ref._ssm_executor.write_bases.copy_(ref._ssm_executor.cursors) accept_unfused(ref, args) if step % 2: graph.replay() @@ -133,8 +146,8 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): checks += 1 class _IndexOnly: - def commit_accepted(self, indices, req_indices, accepted_lengths): - commit_ssm_prefix(indices, req_indices, accepted_lengths) + def commit_accepted(self, indices, req_indices, accepted_lengths, req_verify_width): + indices[req_indices] = (accepted_lengths - 1).to(indices.dtype) m._ssm_executor = _IndexOnly() accept_fused(m, args) @@ -228,14 +241,14 @@ def cycle(i): assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) if iteration % 11 == 10: for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(reqs[:3]) + ssm_backend.fold_accepted_prefix(reqs[:3]) assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) if iteration == 19: for manager in managers: - manager._ssm_executor.clear_history(1) + manager._ssm_executor.cursors[1].zero_() manager._ssm_executor.ssm_state[:, 1].zero_() manager.mtp_state_index[1] = 0 for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(reqs[:3]) + ssm_backend.fold_accepted_prefix(reqs[:3]) assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 0eafbbbb20..6929bf3f9c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -2,12 +2,16 @@ import torch from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, attach_replay_history as ReplaySSMExecutor, ) def ssm_snapshot(backend, req): - return backend.accepted_ssm(req) + state = backend.ssm_state + output = torch.empty((state.shape[0], *state.shape[2:]), dtype=state.dtype, device=state.device) + backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device=state.device), snapshot=output) + return output from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer @@ -79,7 +83,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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) + accept_replay_window(ssm_backend, reqs, accepted, cu) # The correction is larger than FP16 can represent even though the # recurrent state, inputs, and attention outputs are all finite. assert ssm_backend.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max @@ -104,17 +108,16 @@ def test_flush_export_updates_checkpoint_and_clears_history(dtype): backend.betas.fill_(0.5) backend.cursors[0] = 2 expected = ssm_snapshot(backend, 0).clone() - target = torch.empty_like(expected) - - backend.export_accepted_state([0], [target], flush=True) + target = ssm_snapshot(backend, 0) + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + backend.fold_accepted_prefix(reqs) torch.cuda.synchronize() assert backend.cursors[0].item() == 0 torch.testing.assert_close(target, expected, atol=0, rtol=0) torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) - empty_history_target = torch.empty_like(expected) - backend.export_accepted_state([0], [empty_history_target], flush=True) + empty_history_target = ssm_snapshot(backend, 0) torch.cuda.synchronize() torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) @@ -174,7 +177,7 @@ def test_replay_acceptance_flush_and_graph( def step(): out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) return out # Compile on disposable state, then capture with the HOLD slot only. @@ -195,7 +198,7 @@ def step(): for seq, req in enumerate([2, 0]): query_len = width if seq == 0 else max(width - 1, 1) if pending[seq] + query_len > capacity: - ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) + ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) expected[:, req] = expected[:, req].to(dtype).float() pending[seq] = 0 pending[seq] += counts[seq] @@ -223,11 +226,11 @@ def step(): if j + 1 == counts[seq]: expected[layer, req].copy_(cur) if iteration == 2 * capacity + 1: - ssm_backend.merge_accepted_updates(reqs[:2]) + ssm_backend.fold_accepted_prefix(reqs[:2]) expected = expected.to(dtype).float() pending = [0, 0] torch.testing.assert_close(state[:, :4].float(), expected[:, :4], atol=state_atol, rtol=state_rtol) - ssm_backend.merge_accepted_updates(reqs[:2]) + ssm_backend.fold_accepted_prefix(reqs[:2]) torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) @@ -343,10 +346,10 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d ] overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(overflow_reqs) + ssm_backend.fold_accepted_prefix(overflow_reqs) for ssm_backend in ssm_backends: outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) if dtype == torch.float32: torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) else: @@ -355,7 +358,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d 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) + ssm_backend.fold_accepted_prefix(reqs) torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) @@ -370,7 +373,7 @@ def test_bf16_commit_rounds_before_the_next_token(): ssm_backend.gates.fill_(-0.01) ssm_backend.betas.zero_() ssm_backend.cursors[0] = 16 - ssm_backend.merge_accepted_updates(reqs) + ssm_backend.fold_accepted_prefix(reqs) q = torch.ones((1, 1, 32), device="cuda") k = torch.zeros_like(q) v = torch.zeros((1, 1, 64), device="cuda") @@ -423,12 +426,12 @@ def check(capacity, width, steps): # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) for backend in backends: - backend.accept_updates(reqs, accepted, width=width) + accept_replay_window(backend, reqs, accepted, cu) overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] if overflow.numel(): - backend.merge_accepted_updates(overflow) + backend.fold_accepted_prefix(overflow) for backend in backends: - backend.merge_accepted_updates(reqs) + backend.fold_accepted_prefix(reqs) torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) check(4, 1, 16) @@ -447,7 +450,6 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) assert backend.gates.shape[-1] == (dim if kda else capacity) - assert twin.write_bases.shape == backend.cursors.shape reqs = torch.tensor([0], device="cuda", dtype=torch.int32) alog = torch.zeros(hv, device="cuda") bias = torch.zeros(hv * dim if kda else hv, device="cuda") @@ -458,7 +460,7 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() b = torch.randn(1, hv, device="cuda", dtype=dtype) if step == capacity: - twin.merge_accepted_updates(reqs) + twin.fold_accepted_prefix(reqs, run_config=config) assert int(backend.cursors[0]) == capacity outs, twin_outs = [], [] for layer in range(2): @@ -466,11 +468,10 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs)[0]) if step == capacity: assert int(backend.cursors[0]) == capacity - assert int(backend.write_bases[0]) == 0 for actual, expected in zip(outs, twin_outs): torch.testing.assert_close(actual, expected, atol=0, rtol=0) - backend.accept_updates(reqs, reqs, width=1) - twin.accept_updates(reqs, reqs, width=1) + accept_replay_window(backend, reqs) + accept_replay_window(twin, reqs) assert int(backend.cursors[0]) == 1 assert int(twin.cursors[0]) == 1 # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 @@ -494,9 +495,8 @@ def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): b = torch.zeros(1, 2, device="cuda") bias = torch.zeros(2 * 32 if kda else 2, device="cuda") replay_layer(backend, 0, q, k, v, a, b, bias, bias, reqs, cu) - assert int(backend.write_bases[0]) == 3 accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) - backend.accept_updates(reqs, accepted) + accept_replay_window(backend, reqs, accepted, cu) assert backend.cursors.tolist() == [3, 0, 0] @@ -536,7 +536,7 @@ def forward(): def step(): positions = ssm_backend.cursors.index_select(0, reqs.long()) out = forward() - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) return positions, out step() @@ -565,7 +565,7 @@ def step(): before = state.clone() for req, width in zip(ids, widths): if pending[req] + width > 8: - ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) + ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) expected[:, req] = expected[:, req].to(dtype).float() pending[req] = 0 old_cursors = [pending[req] for req in ids] + [0] @@ -573,7 +573,7 @@ def step(): positions.copy_(ssm_backend.cursors.index_select(0, reqs.long())) graph.replay() if not capture_metadata: - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) assert positions.tolist() == old_cursors for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): for layer in range(2): @@ -592,7 +592,7 @@ def step(): if iteration == 2: assert pending[0] == 8 torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) - ssm_backend.merge_accepted_updates(reqs[:2]) + ssm_backend.fold_accepted_prefix(reqs[:2]) torch.testing.assert_close( state.float(), expected.to(dtype).float(), @@ -751,7 +751,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa bias = torch.zeros(32 if kda else 1, device="cuda") for step, query_width in enumerate((4, 4, width)): if step == 2: - backend.clear_history(reqs[folding:live]) + backend.cursors.index_fill_(0, reqs[folding:live].long(), 0) expected[folding:] = initial[0, folding:live] cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 @@ -759,7 +759,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") b = torch.randn(live, query_width, 1, device="cuda") if (backend.cursors[:live] + query_width > backend.capacity).any(): - backend.merge_accepted_updates(reqs[:live]) + backend.fold_accepted_prefix(reqs[:live]) actual = replay_layer( backend, 0, @@ -787,8 +787,9 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa ) torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) - backend.commit_accepted(indices, reqs[:live], lengths) - backend.merge_accepted_updates(reqs[:live]) + backend.advance_cursor(reqs[:live], cu[: live + 1]) + backend.commit_accepted(indices, reqs[:live], lengths, lengths) + backend.fold_accepted_prefix(reqs[:live]) torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) assert torch.equal(backend.ssm_state[0, live], initial[0, live]) @@ -796,7 +797,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): - # 验证窗口放不下时，replay 先折回并把写起点记为 0。提交必须用写起点，不能再把旧 cursor 加上去。 + # 最后一层按整窗全接受推进 cursor。放不下时 cursor 变成窗口长度，而不是旧 cursor 加窗口。 torch.manual_seed(19 + int(kda)) hv, kh, dim = 2, 2, 32 dtype = torch.bfloat16 @@ -817,11 +818,10 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): b = torch.randn(2, hv, device="cuda", dtype=dtype) replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) lengths = torch.tensor([2], device="cuda", dtype=torch.int32) - backend.commit_accepted(indices, reqs, lengths) + backend.advance_cursor(reqs, cu) + backend.commit_accepted(indices, reqs, lengths, lengths) if step < 2: assert int(backend.cursors[0]) == 2 * (step + 1) - assert int(backend.write_bases[0]) == 2 * step else: - assert int(backend.write_bases[0]) == 0 assert int(backend.cursors[0]) == 2 assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 83fe267642..1cde09c8bf 100644 --- 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 @@ -2,6 +2,7 @@ import torch from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, attach_replay_history as ReplaySSMExecutor, ) from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer @@ -47,7 +48,7 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] - backend.accept_updates(ids, ids, width=1) + accept_replay_window(backend, ids) return out graphs = {} @@ -68,8 +69,8 @@ def cycle(backend, ids): for _ in range(8): live = reqs[:count] if (eager.cursors[live.long()] + 1 > eager.capacity).any(): - eager.merge_accepted_updates(live) - captured.merge_accepted_updates(ids[:count]) + eager.fold_accepted_prefix(live) + captured.fold_accepted_prefix(ids[:count]) expected = cycle(eager, reqs[:count]) calls = len(launches) graph.replay() @@ -151,12 +152,17 @@ def forward_accept(index, cu, reqs): ssm_backend = ssm_backends[index] out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] if index == 0: - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) else: + count = starts.numel() + commit_reqs = rows[starts.long()] + window = (cu[1 : count + 1] - cu[:count]).contiguous() + ssm_backend.advance_cursor(commit_reqs, cu[: count + 1]) manager._ssm_executor.commit_accepted( manager.mtp_state_index, - rows[starts.long()], - accepted[rows[starts.long()]] + 1, + commit_reqs, + accepted[commit_reqs.long()] + 1, + window, ) return out @@ -191,7 +197,7 @@ def forward_accept(index, cu, reqs): if overflow: overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(overflow_reqs) + ssm_backend.fold_accepted_prefix(overflow_reqs) folds += len(overflow) for index, ssm_backend in enumerate(ssm_backends): if not capture_metadata: @@ -255,7 +261,7 @@ def forward(ssm_backend): def warm(ssm_backend): out = forward(ssm_backend) - ssm_backend.accept_updates(reqs) + accept_replay_window(ssm_backend, reqs) return out graphs, positions, outputs = [], [], [] @@ -275,7 +281,7 @@ def warm(ssm_backend): pos = ssm_backend.cursors.index_select(0, reqs.long()) out = forward(ssm_backend) if capture_metadata: - ssm_backend.accept_updates(reqs) + accept_replay_window(ssm_backend, reqs) graphs.append(graph) positions.append(pos) outputs.append(out) @@ -287,14 +293,14 @@ def warm(ssm_backend): if overflow: overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(overflow_reqs) + ssm_backend.fold_accepted_prefix(overflow_reqs) old = ssm_backends[0].cursors[reqs.long()].clone() for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): if not capture_metadata: pos.copy_(ssm_backend.cursors.index_select(0, reqs.long())) graph.replay() if not capture_metadata: - ssm_backend.accept_updates(reqs) + accept_replay_window(ssm_backend, reqs) torch.cuda.synchronize() assert torch.equal(positions[0], old) assert torch.equal(positions[0], positions[1]) @@ -338,10 +344,10 @@ def test_cross_layer_fold_with_state_offset_above_int32(kda): def step(ssm_backend, accepted): overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] if overflow: - ssm_backend.merge_accepted_updates(reqs.new_tensor(overflow)) + ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) pos = ssm_backend.cursors.index_select(0, reqs.long()) out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] - ssm_backend.accept_updates(reqs, accepted) + accept_replay_window(ssm_backend, reqs, accepted, cu) return pos, out for count in [2, 3]: @@ -391,7 +397,7 @@ def step(): out = [ replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2) ] - ssm_backend.accept_updates(ids, accepted) + accept_replay_window(ssm_backend, ids, accepted, offsets) return out step() @@ -410,7 +416,7 @@ def step(): if overflow: overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: - ssm_backend.merge_accepted_updates(overflow_reqs) + ssm_backend.fold_accepted_prefix(overflow_reqs) for graph in graphs: graph.replay() for small, padded_out in zip(outputs[0], outputs[1]): @@ -448,16 +454,16 @@ def test_full_prefix_does_not_address_past_capacity(kda): getattr(backend, name).zero_() backend.cursors[reqs.long()] = capacity + 2 replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) - backend.merge_accepted_updates(reqs) + backend.fold_accepted_prefix(reqs) torch.cuda.synchronize() assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) assert torch.isfinite(backend.ssm_state).all() for _ in range(12): replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) - backend.accept_updates(reqs, reqs, width=1) + accept_replay_window(backend, reqs) overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] if overflow.numel(): - backend.merge_accepted_updates(overflow) + backend.fold_accepted_prefix(overflow) torch.cuda.synchronize() assert int(backend.cursors[reqs.long()].max()) <= capacity assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index a4df60ae25..f16d81788d 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -165,9 +165,8 @@ def test_native_commit_selects_mtp_slot(cls): if not torch.cuda.is_available(): pytest.skip("CUDA required") owner, _, _ = make_state(cls, step=2) - owner.commit_bounded_att_state( - torch.tensor([1], device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") - ) + accepted = torch.tensor([3], dtype=torch.int32, device="cuda") + owner.commit_bounded_att_state(torch.tensor([1], device="cuda"), accepted, accepted) assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] @@ -179,13 +178,27 @@ def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): history = getattr(owner, "runtime", owner) calls = [] monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) - monkeypatch.setattr(history, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) - owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3])) + widths = torch.tensor([1, 2, 3]) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), widths) assert len(calls) == 1 - assert calls[0][3] is history.cursors + assert calls[0][3] is widths + assert calls[0][4] is history.cursors assert calls[0][0] is owner.mtp_state_index +def _commit_prefix_cpu(indices, reqs, lengths, req_verify_width, replay_cursors) -> None: + """CPU 上按请求行核对索引和 cursor。正式提交算子只跑 CUDA。""" + + indices[reqs] = (lengths - 1).to(indices.dtype) + hold = replay_cursors.numel() - 1 + for row in range(reqs.numel()): + req = int(reqs[row]) + if req == hold: + continue + rejected = int(req_verify_width[row]) - int(lengths[row]) + replay_cursors[req] = int(replay_cursors[req]) - rejected + + @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_acceptance_keeps_request_order_with_strided_inputs(device): if device == "cuda" and not torch.cuda.is_available(): @@ -193,10 +206,19 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): owner, _, _ = make_state(mode="replay", step=3, device=device) owner.mtp_state_index.fill_(2) history = getattr(owner, "runtime", owner) - history.cursors.fill_(8) - history.write_bases.copy_(history.cursors) + history.cursors.zero_() reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] - owner.commit_accepted(owner.mtp_state_index, reqs, lengths) + if device == "cuda": + flat_lengths = lengths.contiguous() + cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) + cu[1:] = torch.cumsum(flat_lengths, 0) + history.advance_cursor(reqs.contiguous(), cu) + owner.commit_accepted(owner.mtp_state_index, reqs, lengths, lengths) + else: + # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 + history.cursors[0] = 3 + history.cursors[2] = 1 + _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] - assert history.cursors.tolist() == [11, 8, 9, 8] + assert history.cursors.tolist() == [3, 0, 1, 0] diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index bc6c1428cb..31b6236139 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -289,14 +289,14 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod 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 - selected.accept_updates(ids, accepted) for req, length in zip(requests, lengths): + snap = torch.empty_like(selected.ssm_state[:, req]) + selected.fold_accepted_prefix( + torch.tensor([req], dtype=torch.int32, device="cuda"), + snapshot=snap, + ) torch.testing.assert_close( - selected.accepted_ssm(req)[0], + snap[0], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, From 69c2518ec9058f9f1a72837b9fabf4408c095332 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 13:50:06 +0000 Subject: [PATCH 37/54] fix: wait for every value tile before a folded replay writes history Folding rewrites shared raw keys and gates at slot 0. Each launch row now waits until its value tiles have finished reading the old prefix, so a late tile cannot fold the new token into the SSM. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 2 ++ .../bounded_att/kda/replay_kda_state.py | 2 ++ .../attention/linear/replay_gdn_decode.py | 3 +++ .../attention/linear/replay_kda_decode.py | 3 +++ .../linear_att/replayssm_decode.py | 20 +++++++++++++++++++ .../linear_att/replay_history.py | 1 + 6 files changed, 31 insertions(+) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 9506fc6d4e..5796574c93 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -103,6 +103,8 @@ def __init__( self.hold_req_id = slots - 1 # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # GDN 的 g 是每个头一个标量。 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index daa297dbb9..28f3e54932 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -103,6 +103,8 @@ def __init__( self.hold_req_id = slots - 1 # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # KDA 的 g 按 key 通道存。 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 8bfbcf07ea..6d0f1ca405 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -44,6 +44,7 @@ def replay( precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) + sync_mark = history.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -61,6 +62,7 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, + sync_mark, cu_seqlens, out, q.stride(axis), @@ -84,6 +86,7 @@ def replay( history.gate_lower_bound, BQK=max(16, 2 * bw), BR=triton.next_power_of_2(max(16, history.capacity + bw)), + NV=triton.cdiv(vd, bv), num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index e9f8c2072f..0c72a28919 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -44,6 +44,7 @@ def replay( precompute_projections = history.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 bw = triton.next_power_of_2(history.verify_width) + sync_mark = history.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -61,6 +62,7 @@ def replay( history.betas[linear_layer_idx], reqs, history.cursors, + sync_mark, cu_seqlens, out, q.stride(axis), @@ -84,6 +86,7 @@ def replay( history.gate_lower_bound, BQK=max(16, 2 * bw), BR=triton.next_power_of_2(max(16, history.capacity + bw)), + NV=triton.cdiv(vd, bv), num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index fae106be22..8bee9f5251 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -12,6 +12,22 @@ from .replayssm_delta import _apply_accepted_prefix +@triton.jit +def _arrive_value_tiles(SyncMark, seq, hv, HV: tl.constexpr, NV: tl.constexpr): + # 同一个 launch 行、同一个 value head 的 V/BV 个 tile 都读完旧前缀后，才允许写新窗口。 + # 行号不是请求槽。HOLD 会在一批里重复出现，用请求槽做下标会撞到同一格。 + # 最后到达的 tile 把计数写回 0，这个 0 就是放行信号。其余 tile 已经加过 1， + # 之后再看见 0 只可能是这次复原。 + idx = seq * HV + hv + arrived = tl.atomic_add(SyncMark + idx, 1, sem="acq_rel", scope="gpu") + if arrived == NV - 1: + tl.atomic_xchg(SyncMark + idx, 0, sem="release", scope="gpu") + else: + seen = arrived + 1 + while seen != 0: + seen = tl.atomic_add(SyncMark + idx, 0, sem="acquire", scope="gpu") + + @triton.jit def _history_coeff(rows, vec, BK: tl.constexpr): # rows [R, BK] · vec [BK] -> [R]. tf32x3 是 tensor core 上最接近 fp32 的输出重建。 @@ -46,6 +62,7 @@ def _decode_replay( Betas, Reqs, Cursors, + SyncMark, Cu, Out, SQ: tl.constexpr, @@ -69,6 +86,7 @@ def _decode_replay( LOWER: tl.constexpr, BQK: tl.constexpr, BR: tl.constexpr, + NV: tl.constexpr, ): """执行一个请求窗口的 replay decode，并写入 history。 @@ -103,6 +121,8 @@ def _decode_replay( state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) n = 0 + # 新 token 会从槽 0 覆盖刚读过的 RawKeys、Gates、Betas。等这一组 V tile 都读完再往下走。 + _arrive_value_tiles(SyncMark, seq, hv, HV, NV) if PRECOMPUTE_PROJECTIONS: # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index ede0beae91..f0f94c8412 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -56,6 +56,7 @@ def attach_replay_history( # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度。 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.tile_sync_mark = torch.zeros((slots * (owner.mtp_step + 6), hv), dtype=torch.int32, device=ssm.device) # 重建用 k/d：bf16 激活用 fp16。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) From 5dc284e769570480b7155e47d42e6799d27944d5 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 14:04:10 +0000 Subject: [PATCH 38/54] refactor: require cursor prefixes at decode setup Width-1 decode now builds its cumulative lengths while initializing indices. The last layer advances the cursor with that tensor instead of allocating one when the argument is missing. Co-authored-by: Cursor --- .../attention_state/bounded_att/gdn/replay_gdn_state.py | 6 +----- .../attention_state/bounded_att/kda/replay_kda_state.py | 6 +----- lightllm/common/basemodel/attention/linear/replay_gdn.py | 9 ++++++++- lightllm/common/basemodel/attention/linear/replay_kda.py | 9 ++++++++- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 5796574c93..4a2b9af2c2 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -426,16 +426,12 @@ def read_from_pd_page( # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 # ======================================================================== - def finish_decode_token( - self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None - ) -> None: + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - if cu_seqlens is None: - cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) self.advance_cursor(reqs, cu_seqlens) def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 28f3e54932..4ea3565f5d 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -426,16 +426,12 @@ def read_from_pd_page( # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 # ======================================================================== - def finish_decode_token( - self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None - ) -> None: + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - if cu_seqlens is None: - cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) self.advance_cursor(reqs, cu_seqlens) def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index eb54adcfaa..3892648934 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -154,6 +154,7 @@ class ReplayGDNDecodeAttState(BaseDecodeAttState): b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None + b_cu_seqlens: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None def init_state(self): @@ -191,7 +192,7 @@ def decode_att( else: core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) return core, z def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: @@ -200,6 +201,10 @@ def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None self.b_conv_buffer_idx = infer.b_req_idx self.b1_mtp_cu_q_seq_len = None self.b_num_accepted_tokens = None + # 每行一个 token。replay kernel 仍用这个布局，cursor 需要现成的前缀和。 + self.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) return if self.backend.uses_dynamic_spec_verify_layout(): cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( @@ -209,6 +214,7 @@ def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None hold_req_id=infer.req_manager.HOLD_REQUEST_ID, ) self.b1_mtp_cu_q_seq_len = cu + self.b_cu_seqlens = cu self.b_conv_buffer_idx = conv_idx self.b_num_accepted_tokens = accepted return @@ -222,6 +228,7 @@ def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None self.b1_mtp_cu_q_seq_len = torch.arange( 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device ) + self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index a03f167563..0354cfe5fc 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -131,6 +131,7 @@ class ReplayKDADecodeAttState(BaseDecodeAttState): b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None + b_cu_seqlens: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None def init_state(self): @@ -196,7 +197,7 @@ def decode_att( self.b1_mtp_cu_q_seq_len, ) # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) return output def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: @@ -205,6 +206,10 @@ def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: self.b_conv_buffer_idx = infer.b_req_idx self.b1_mtp_cu_q_seq_len = None self.b_num_accepted_tokens = None + # 每行一个 token。replay kernel 仍用这个布局，cursor 需要现成的前缀和。 + self.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) return if self.backend.uses_dynamic_spec_verify_layout(): cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( @@ -214,6 +219,7 @@ def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: hold_req_id=infer.req_manager.HOLD_REQUEST_ID, ) self.b1_mtp_cu_q_seq_len = cu + self.b_cu_seqlens = cu self.b_conv_buffer_idx = conv_idx self.b_num_accepted_tokens = accepted return @@ -227,6 +233,7 @@ def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: self.b1_mtp_cu_q_seq_len = torch.arange( 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device ) + self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 From 1d9162a4c7e13bccc0ebb50ff78e441d9c96505e Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 14:22:22 +0000 Subject: [PATCH 39/54] refactor: keep one MTP index path and commit only when MTP is enabled The index tensor is always allocated, and verification trims the replay cursor through commit_ssm_prefix. Co-authored-by: Cursor --- .../bounded_att/gdn/gdn_state.py | 13 ++++--- .../bounded_att/gdn/replay_gdn_state.py | 37 +++++-------------- .../bounded_att/kda/kda_state.py | 13 ++++--- .../bounded_att/kda/replay_kda_state.py | 37 +++++-------------- .../test_model_initialization.py | 4 +- .../linear_att/replay_history.py | 5 ++- .../linear_att/test_acceptance_fusion.py | 17 +++------ .../linear_att/test_replayssm.py | 5 ++- .../linear_att/test_replayssm_execution.py | 4 +- unit_tests/common/test_linear_state.py | 8 ++-- .../common/test_linear_state_contract.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 4 +- 12 files changed, 57 insertions(+), 92 deletions(-) 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 b001d74ae3..82eb4026c5 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -46,6 +46,7 @@ def __init__( self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) @@ -67,7 +68,7 @@ def __init__( dtype=config.ssm_dtype, device=device, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None @@ -126,7 +127,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: self.conv_state[:, req_idx].zero_() start = req_idx * (self.mtp_step + 1) self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -170,7 +171,7 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -182,7 +183,7 @@ def commit_bounded_att_state( ) -> None: """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_indices] = accepted_steps - 1 # ======================================================================== @@ -321,7 +322,7 @@ def copy_runtime_state_from_dp( ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -376,7 +377,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 4a2b9af2c2..bed3a6d86c 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -51,6 +51,7 @@ def __init__( self.state_name = None self.request_capacity = request_capacity self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self.gate_lower_bound = gate_lower_bound self._algorithm = "gdn" self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} @@ -76,7 +77,7 @@ def __init__( dtype=config.ssm_dtype, device=device, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None replay_config = dict(ssm_replay_config or {}) @@ -179,7 +180,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() self.cursors[req_idx].zero_() - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -219,7 +220,7 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -229,10 +230,10 @@ def commit_bounded_att_state( accepted_steps: torch.Tensor, req_verify_width: torch.Tensor, ) -> None: - """提交已接受长度，并剪掉整窗全接受后未通过验证的尾部。""" - if self.mtp_state_index is None: + """MTP 验证后写入索引，并从整窗全接受后的 cursor 上剪掉未接受尾部。""" + if not self.mtp_enabled: return - self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps, req_verify_width) + commit_ssm_prefix(self.mtp_state_index, req_indices, accepted_steps, req_verify_width, self.cursors) # ======================================================================== # 三、CPU Cache 数据交换 @@ -353,7 +354,7 @@ def copy_runtime_state_from_dp( ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -414,7 +415,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== @@ -439,26 +440,6 @@ def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: advance_replay_cursor(self, reqs, cu_seqlens) - def commit_accepted( - self, - mtp_state_index: torch.Tensor, - reqs: torch.Tensor, - accepted_lengths: torch.Tensor, - req_verify_width: torch.Tensor, - ) -> None: - """MTP verify 之后写入索引，并从全接受后的 cursor 上剪掉未接受尾部。 - - accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 - 配置的 verify_width 为 1 时没有尾部，只把这一个 token 记入 cursor。 - """ - - if self.verify_width > 1: - commit_ssm_prefix(mtp_state_index, reqs, accepted_lengths, req_verify_width, self.cursors) - else: - mtp_state_index[reqs] = (accepted_lengths - 1).to(mtp_state_index.dtype) - cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - self.advance_cursor(reqs, cu_seqlens) - def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 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 9e386dbe2d..68097905d5 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -46,6 +46,7 @@ def __init__( self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) @@ -67,7 +68,7 @@ def __init__( dtype=config.ssm_dtype, device=device, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None @@ -126,7 +127,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: self.conv_state[:, req_idx].zero_() start = req_idx * (self.mtp_step + 1) self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -170,7 +171,7 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -182,7 +183,7 @@ def commit_bounded_att_state( ) -> None: """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_indices] = accepted_steps - 1 # ======================================================================== @@ -321,7 +322,7 @@ def copy_runtime_state_from_dp( ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -376,7 +377,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 4ea3565f5d..67a2abcfe0 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -51,6 +51,7 @@ def __init__( self.state_name = None self.request_capacity = request_capacity self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self.gate_lower_bound = gate_lower_bound self._algorithm = "kda" self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} @@ -76,7 +77,7 @@ def __init__( dtype=config.ssm_dtype, device=device, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None replay_config = dict(ssm_replay_config or {}) @@ -179,7 +180,7 @@ def reset_bounded_att_state(self, req_indices: list[int]) -> None: self.conv_state[:, req_idx].zero_() self.ssm_state[:, req_idx].zero_() self.cursors[req_idx].zero_() - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -219,7 +220,7 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -229,10 +230,10 @@ def commit_bounded_att_state( accepted_steps: torch.Tensor, req_verify_width: torch.Tensor, ) -> None: - """提交已接受长度，并剪掉整窗全接受后未通过验证的尾部。""" - if self.mtp_state_index is None: + """MTP 验证后写入索引，并从整窗全接受后的 cursor 上剪掉未接受尾部。""" + if not self.mtp_enabled: return - self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps, req_verify_width) + commit_ssm_prefix(self.mtp_state_index, req_indices, accepted_steps, req_verify_width, self.cursors) # ======================================================================== # 三、CPU Cache 数据交换 @@ -353,7 +354,7 @@ def copy_runtime_state_from_dp( ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -414,7 +415,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== @@ -439,26 +440,6 @@ def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: advance_replay_cursor(self, reqs, cu_seqlens) - def commit_accepted( - self, - mtp_state_index: torch.Tensor, - reqs: torch.Tensor, - accepted_lengths: torch.Tensor, - req_verify_width: torch.Tensor, - ) -> None: - """MTP verify 之后写入索引，并从全接受后的 cursor 上剪掉未接受尾部。 - - accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 - 配置的 verify_width 为 1 时没有尾部，只把这一个 token 记入 cursor。 - """ - - if self.verify_width > 1: - commit_ssm_prefix(mtp_state_index, reqs, accepted_lengths, req_verify_width, self.cursors) - else: - mtp_state_index[reqs] = (accepted_lengths - 1).to(mtp_state_index.dtype) - cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - self.advance_cursor(reqs, cu_seqlens) - def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index af8dbfcde0..09e88b4a93 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -162,7 +162,9 @@ def check_runtime_before_profile(device: str | torch.device | int): assert resources.bounded_att_big_page_allocator.size == 6 assert resources.bounded_att_small_page_allocator.size == 2 assert not hasattr(model, "mtp_state_cursors") - assert (main_linear_state(resources).mtp_state_index is not None) == bool(mtp) + linear_state = main_linear_state(resources) + assert linear_state.mtp_state_index is not None + assert linear_state.mtp_enabled == bool(mtp) if mode == "prefill": assert resources.get_pd_operator().get_pd_full_att_page_layouts() assert resources.get_pd_operator().get_pd_bounded_att_page_layouts() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index f0f94c8412..95c24af2fe 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -35,10 +35,11 @@ def attach_replay_history( owner.state_name = None owner.request_capacity = ssm.shape[1] owner.mtp_step = verify_width - 1 + owner.mtp_enabled = owner.mtp_step > 0 owner._algorithm = "kda" if kda else "gdn" owner._logical_to_storage_layer = {} owner.conv_state = None - owner.mtp_state_index = None + owner.mtp_state_index = torch.zeros(owner.request_capacity, dtype=torch.int32, device=ssm.device) owner.big_pages_pool = None owner.small_pages_pool = None owner.ssm_state = ssm @@ -93,7 +94,7 @@ def _trim_rejected_index( def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens=None) -> None: """测试里一次结算 cursor：先按整窗全接受，再按槽位 index 剪掉未接受尾部。 - 没传 cu_seqlens 时每个请求的窗口长度是 1。正式路径用 advance_cursor 和 commit_accepted。 + 没传 cu_seqlens 时每个请求的窗口长度是 1。正式路径先 advance_cursor，开了 MTP 再 commit_ssm_prefix。 """ if reqs.numel() == 0: 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 index d384fadd92..53797874a5 100644 --- 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 @@ -1,6 +1,7 @@ import pytest import torch +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, @@ -97,12 +98,14 @@ def accept_fused(m, args): cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) lengths = cumulative[ends.long()] - cumulative[args[0].long()] executor = m._ssm_executor - if getattr(executor, "verify_width", 1) > 1: + if executor.verify_width > 1: cu = _replay_cu(args) executor.advance_cursor(reqs, cu) - executor.commit_accepted(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1]) + commit_ssm_prefix(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1], executor.cursors) else: - executor.commit_accepted(m.mtp_state_index, reqs, lengths, lengths) + # 宽度为 1 时正式路径不提交。对照只写下标，并把每个请求的 cursor 加一。 + m.mtp_state_index[reqs] = (lengths - 1).to(m.mtp_state_index.dtype) + executor.advance_cursor(reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -145,14 +148,6 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): assert m._ssm_executor.cursors.tolist() == expected_cursor checks += 1 - class _IndexOnly: - def commit_accepted(self, indices, req_indices, accepted_lengths, req_verify_width): - indices[req_indices] = (accepted_lengths - 1).to(indices.dtype) - - m._ssm_executor = _IndexOnly() - accept_fused(m, args) - assert m.mtp_state_index.tolist() == expected_index - print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) 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 index 6929bf3f9c..505e1340f0 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -14,6 +14,7 @@ def ssm_snapshot(backend, req): return output +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -788,7 +789,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) backend.advance_cursor(reqs[:live], cu[: live + 1]) - backend.commit_accepted(indices, reqs[:live], lengths, lengths) + commit_ssm_prefix(indices, reqs[:live], lengths, lengths, backend.cursors) backend.fold_accepted_prefix(reqs[:live]) torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) assert torch.equal(backend.ssm_state[0, live], initial[0, live]) @@ -819,7 +820,7 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) lengths = torch.tensor([2], device="cuda", dtype=torch.int32) backend.advance_cursor(reqs, cu) - backend.commit_accepted(indices, reqs, lengths, lengths) + commit_ssm_prefix(indices, reqs, lengths, lengths, backend.cursors) if step < 2: assert int(backend.cursors[0]) == 2 * (step + 1) else: 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 index 1cde09c8bf..3132c5ecb8 100644 --- 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 @@ -9,6 +9,7 @@ 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.ssm_acceptance import commit_ssm_prefix from types import SimpleNamespace pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -158,11 +159,12 @@ def forward_accept(index, cu, reqs): commit_reqs = rows[starts.long()] window = (cu[1 : count + 1] - cu[:count]).contiguous() ssm_backend.advance_cursor(commit_reqs, cu[: count + 1]) - manager._ssm_executor.commit_accepted( + commit_ssm_prefix( manager.mtp_state_index, commit_reqs, accepted[commit_reqs.long()] + 1, window, + ssm_backend.cursors, ) return out diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 02dfe2a6ca..bb2314a025 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -47,7 +47,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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: + if state.mtp_enabled: state.mtp_state_index[src] = offset if mode == "native": expected_conv = state.conv_state[:, src, ..., :3].clone() @@ -80,7 +80,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( 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: + if state.mtp_enabled: state.mtp_state_index[dst] = step if mode == "replay": state.cursors[dst] = 3 @@ -92,7 +92,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) - if state.mtp_state_index is not None: + if state.mtp_enabled: assert state.mtp_state_index[dst].item() == 0 if mode == "replay": assert state.cursors[dst].item() == 0 @@ -156,7 +156,7 @@ def test_checkpoint_batch_export(runtime_start_args, cls, mode, dtype, destinati if cursors is not None: assert torch.equal(state.cursors, cursors) target, _, _ = make_state(cls, mode, 2, dtype) - if target.mtp_state_index is not None: + if target.mtp_enabled: target.mtp_state_index.fill_(2) target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) for req, (conv, ssm) in enumerate(expected): diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index f16d81788d..cf4eb592d8 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -214,7 +214,7 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) cu[1:] = torch.cumsum(flat_lengths, 0) history.advance_cursor(reqs.contiguous(), cu) - owner.commit_accepted(owner.mtp_state_index, reqs, lengths, lengths) + owner.commit_bounded_att_state(reqs, lengths, lengths) else: # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 history.cursors[0] = 3 diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index d230d58839..4d15d6bf1c 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -189,7 +189,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.ssm_state.normal_() if dest_mode == "replay": mem.linear.cursors[dst_req] = 3 - if mem.linear.mtp_state_index is not None: + if mem.linear.mtp_enabled: mem.linear.mtp_state_index[dst_req] = decode_mtp_step for state in mem.tails: state.buffer.normal_() @@ -261,7 +261,7 @@ def test_pd_kv_and_runtime_state_roundtrip( assert torch.equal(restored_tail, expected_tail) if dest_mode == "replay": assert mem.linear.cursors[dst_req] == 0 - if mem.linear.mtp_state_index is not None: + if mem.linear.mtp_enabled: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( ( From 64b6f1f148934e2ac5768f04fd03f415d669c59c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 14:30:29 +0000 Subject: [PATCH 40/54] refactor: fold replay prefixes only back into the canonical SSM The state no longer wraps cursor updates or exposes a snapshot fold. Tests that need a read-only result fold a clone. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 27 ++++----------- .../bounded_att/kda/replay_kda_state.py | 27 ++++----------- .../linear_att/replayssm_fold.py | 34 +++---------------- .../triton_kernel/linear_att/ssm_autotune.py | 3 +- .../linear_att/replay_history.py | 3 +- .../linear_att/test_acceptance_fusion.py | 20 +++++++---- .../linear_att/test_replayssm.py | 23 +++++++++---- .../linear_att/test_replayssm_execution.py | 3 +- .../common/test_linear_state_contract.py | 3 +- unit_tests/models/glm5_next/test_mtp.py | 15 ++++---- 10 files changed, 66 insertions(+), 92 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index bed3a6d86c..8508123d50 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -11,7 +11,7 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor as advance_replay_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -424,7 +424,7 @@ def read_from_pd_page( # cursor 是已接受、尚未折回的长度，各层只读。 # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 - # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 # ======================================================================== def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: @@ -433,28 +433,15 @@ def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seql layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - self.advance_cursor(reqs, cu_seqlens) + advance_cursor(self, reqs, cu_seqlens) - def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: - """按 cu_seqlens 给出的整窗全接受推进 cursor，不修改 SSM。""" - - advance_replay_cursor(self, reqs, cu_seqlens) - - def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: - """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 - - snapshot 为空时写回 canonical SSM，并清掉这些请求的 cursor。 - 传入 snapshot 时只写入这块独立缓冲，不改 canonical SSM，也不动 cursor。 - run_config 为空时由 fold_prefix 使用它自己的默认配置。 - """ + def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 canonical SSM，并清掉这些请求的 cursor。""" if reqs.numel() == 0: return - if snapshot is None: - fold_prefix(self, reqs, self.ssm_state, False, run_config) - self.cursors.index_fill_(0, reqs.long(), 0) - return - fold_prefix(self, reqs, snapshot, True, run_config) + fold_prefix(self, reqs) + self.cursors.index_fill_(0, reqs.long(), 0) # ======================================================================== # 七、私有辅助实现 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 67a2abcfe0..983eccb5c8 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -11,7 +11,7 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor as advance_replay_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger @@ -424,7 +424,7 @@ def read_from_pd_page( # cursor 是已接受、尚未折回的长度，各层只读。 # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 - # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 只用于层外对照和只读快照。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 # ======================================================================== def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: @@ -433,28 +433,15 @@ def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seql layers = () if self.config is None else self.config.logical_layer_ids if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: return - self.advance_cursor(reqs, cu_seqlens) + advance_cursor(self, reqs, cu_seqlens) - def advance_cursor(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: - """按 cu_seqlens 给出的整窗全接受推进 cursor，不修改 SSM。""" - - advance_replay_cursor(self, reqs, cu_seqlens) - - def fold_accepted_prefix(self, reqs: torch.Tensor, snapshot: torch.Tensor | None = None, run_config=None) -> None: - """用原始 k/v 把 cursor 覆盖的已接受前缀折进 SSM。 - - snapshot 为空时写回 canonical SSM，并清掉这些请求的 cursor。 - 传入 snapshot 时只写入这块独立缓冲，不改 canonical SSM，也不动 cursor。 - run_config 为空时由 fold_prefix 使用它自己的默认配置。 - """ + def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 canonical SSM，并清掉这些请求的 cursor。""" if reqs.numel() == 0: return - if snapshot is None: - fold_prefix(self, reqs, self.ssm_state, False, run_config) - self.cursors.index_fill_(0, reqs.long(), 0) - return - fold_prefix(self, reqs, snapshot, True, run_config) + fold_prefix(self, reqs) + self.cursors.index_fill_(0, reqs.long(), 0) # ======================================================================== # 七、私有辅助实现 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index b0a9e48dda..0d1af0618b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -10,7 +10,6 @@ @triton.jit def _fold_prefix( State, - Output, RawKeys, RawValues, Gates, @@ -26,23 +25,20 @@ def _fold_prefix( HOLD: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, - SNAPSHOT: tl.constexpr, KDA: tl.constexpr, - FLUSH_EXPORT: tl.constexpr = False, - OUTPUT_STRIDES: tl.constexpr = None, STRIDE_LAYER: tl.constexpr = 0, STRIDE_SLOT: tl.constexpr = 0, STRIDE_HEAD: tl.constexpr = 0, STRIDE_K: tl.constexpr = 0, STRIDE_V: tl.constexpr = 1, ): - # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT 只写独立输出。 + # 将 cursor 指向的 accepted history 写回 canonical state。 iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) req = tl.load(Reqs + row).to(tl.int64) if req == HOLD: return n = tl.load(Cursors + req) - if n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + if n == 0: return linear_layer_idx, head = lh // HV, lh % HV # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 @@ -67,31 +63,13 @@ def _fold_prefix( state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) - mask = (kk[:, None] < K) & (vv[None, :] < V) - if OUTPUT_STRIDES is not None: - target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) - out = ( - target - + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] - + head * OUTPUT_STRIDES[1] - + kk[:, None] * OUTPUT_STRIDES[2] - + vv[None, :] * OUTPUT_STRIDES[3] - ) - elif SNAPSHOT: - out_slot = layer64 * HV + head64 - out = Output + out_slot * K * V + kk[:, None].to(tl.int64) * V + vv[None, :].to(tl.int64) - else: - out = sp - tl.store(out, state, mask) - if FLUSH_EXPORT: - tl.store(sp, state, mask) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) # cursor 不在这里清。和折回并行的其他线程块还要读原来的长度，清零放到 kernel 结束之后。 -def fold_prefix(history, reqs, output, snapshot: bool, run_config=None) -> None: - """把 history 上 cursor 覆盖的已接受前缀折进 output。 +def fold_prefix(history, reqs, run_config=None) -> None: + """把 history 上 cursor 覆盖的已接受前缀折进 canonical SSM。 - snapshot 为真时 output 是独立缓冲，kernel 不写回 canonical SSM。 没传入 run_config 时用 BV=32、4 个 warp、3 个 stage。 """ @@ -105,7 +83,6 @@ def fold_prefix(history, reqs, output, snapshot: bool, run_config=None) -> None: layer_stride, slot_stride, head_stride, key_stride, value_stride = state.stride() _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( state, - output, history.raw_keys, history.raw_values, history.gates, @@ -121,7 +98,6 @@ def fold_prefix(history, reqs, output, snapshot: bool, run_config=None) -> None: history.hold_req_id, triton.next_power_of_2(k), bv, - snapshot, history.kda, STRIDE_LAYER=layer_stride, STRIDE_SLOT=slot_stride, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 59b1ff5583..ce761a3483 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -7,6 +7,7 @@ import torch import triton +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune from lightllm.utils.log_utils import init_logger @@ -166,7 +167,7 @@ def select_config( adopt_cu = cu_seqlens if adopt_cu is None: adopt_cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - ssm_backend.advance_cursor(reqs, adopt_cu) + advance_cursor(ssm_backend, reqs, adopt_cu) return run_config diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 95c24af2fe..84b6e8db99 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -6,6 +6,7 @@ from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor def attach_replay_history( @@ -101,7 +102,7 @@ def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens= return if cu_seqlens is None: cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - history.advance_cursor(reqs, cu_seqlens) + advance_cursor(history, reqs, cu_seqlens) if cu_seqlens is None or accepted_state_indices is None: return _trim_rejected_index[(triton.cdiv(reqs.numel(), 256),)]( 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 index 53797874a5..f6cac54b6a 100644 --- 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 @@ -1,6 +1,8 @@ import pytest import torch +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( @@ -10,10 +12,16 @@ def ssm_snapshot(backend, req): - state = backend.ssm_state - output = torch.empty((state.shape[0], *state.shape[2:]), dtype=state.dtype, device=state.device) - backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device=state.device), snapshot=output) - return output + """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" + + saved = backend.ssm_state + clone = saved.clone() + backend.ssm_state = clone + try: + fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) + finally: + backend.ssm_state = saved + return clone[:, req] from types import SimpleNamespace @@ -100,12 +108,12 @@ def accept_fused(m, args): executor = m._ssm_executor if executor.verify_width > 1: cu = _replay_cu(args) - executor.advance_cursor(reqs, cu) + advance_cursor(executor, reqs, cu) commit_ssm_prefix(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1], executor.cursors) else: # 宽度为 1 时正式路径不提交。对照只写下标，并把每个请求的 cursor 加一。 m.mtp_state_index[reqs] = (lengths - 1).to(m.mtp_state_index.dtype) - executor.advance_cursor(reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) + advance_cursor(executor, reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") 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 index 505e1340f0..53f38e15d1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -8,12 +8,20 @@ def ssm_snapshot(backend, req): - state = backend.ssm_state - output = torch.empty((state.shape[0], *state.shape[2:]), dtype=state.dtype, device=state.device) - backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device=state.device), snapshot=output) - return output + """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" + saved = backend.ssm_state + clone = saved.clone() + backend.ssm_state = clone + try: + fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) + finally: + backend.ssm_state = saved + return clone[:, req] + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer @@ -461,7 +469,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() b = torch.randn(1, hv, device="cuda", dtype=dtype) if step == capacity: - twin.fold_accepted_prefix(reqs, run_config=config) + fold_prefix(twin, reqs, config) + twin.cursors.index_fill_(0, reqs.long(), 0) assert int(backend.cursors[0]) == capacity outs, twin_outs = [], [] for layer in range(2): @@ -788,7 +797,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa ) torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) - backend.advance_cursor(reqs[:live], cu[: live + 1]) + advance_cursor(backend, reqs[:live], cu[: live + 1]) commit_ssm_prefix(indices, reqs[:live], lengths, lengths, backend.cursors) backend.fold_accepted_prefix(reqs[:live]) torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) @@ -819,7 +828,7 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): b = torch.randn(2, hv, device="cuda", dtype=dtype) replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) lengths = torch.tensor([2], device="cuda", dtype=torch.int32) - backend.advance_cursor(reqs, cu) + advance_cursor(backend, reqs, cu) commit_ssm_prefix(indices, reqs, lengths, lengths, backend.cursors) if step < 2: assert int(backend.cursors[0]) == 2 * (step + 1) 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 index 3132c5ecb8..28bfdea54d 100644 --- 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 @@ -9,6 +9,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix from types import SimpleNamespace @@ -158,7 +159,7 @@ def forward_accept(index, cu, reqs): count = starts.numel() commit_reqs = rows[starts.long()] window = (cu[1 : count + 1] - cu[:count]).contiguous() - ssm_backend.advance_cursor(commit_reqs, cu[: count + 1]) + advance_cursor(ssm_backend, commit_reqs, cu[: count + 1]) commit_ssm_prefix( manager.mtp_state_index, commit_reqs, diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index cf4eb592d8..f27909f889 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -4,6 +4,7 @@ from lightllm.common.attention_state.bounded_att import KDAState from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState @@ -213,7 +214,7 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): flat_lengths = lengths.contiguous() cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) cu[1:] = torch.cumsum(flat_lengths, 0) - history.advance_cursor(reqs.contiguous(), cu) + advance_cursor(history, reqs.contiguous(), cu) owner.commit_bounded_att_state(reqs, lengths, lengths) else: # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 31b6236139..30c4fc22c4 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -8,6 +8,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( fused_recurrent_kda, ) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -290,13 +291,15 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod ) else: for req, length in zip(requests, lengths): - snap = torch.empty_like(selected.ssm_state[:, req]) - selected.fold_accepted_prefix( - torch.tensor([req], dtype=torch.int32, device="cuda"), - snapshot=snap, - ) + saved = selected.ssm_state + clone = saved.clone() + selected.ssm_state = clone + try: + fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) + finally: + selected.ssm_state = saved torch.testing.assert_close( - snap[0], + clone[0, req], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, From e97fce7df71195ba2e93ec9c53c8d2939a43c707 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 11 Nov 2026 14:47:23 +0000 Subject: [PATCH 41/54] refactor: drop unused replay kernel exports and the dense stride fallback Decode calls the replay kernel directly, and the external fold addresses canonical SSM with its real strides. Co-authored-by: Cursor --- .../attention/linear/replay_gdn_decode.py | 2 +- .../attention/linear/replay_kda_decode.py | 2 +- .../triton_kernel/linear_att/replayssm.py | 30 ------------------- .../linear_att/replayssm_decode.py | 22 +++----------- .../linear_att/replayssm_fold.py | 22 ++++++-------- .../linear_att/test_replayssm_execution.py | 8 ++--- 6 files changed, 19 insertions(+), 67 deletions(-) delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 6d0f1ca405..cc7b12b248 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -8,7 +8,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 0c72a28919..8c2684140c 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -8,7 +8,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py deleted file mode 100644 index 7b57371163..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""ReplaySSM kernel exports, in lifecycle order. - -decode writes history for the current window and does not move the cursor. -advance_cursor moves the cursor after the last layer; verify trims the rejected tail. -fold applies that accepted prefix onto the canonical SSM. -""" - -from .replayssm_decode import ( - _decode_replay, - _gdn_precomputed_window, - _kda_precomputed_window, - _project_qk, - _replay_each_token, -) -from .replayssm_delta import _apply_accepted_prefix, _delta_rule_step -from .replayssm_cursor import _advance_cursor -from .replayssm_fold import _fold_prefix - -__all__ = [ - "_decode_replay", - "_project_qk", - "_gdn_precomputed_window", - "_kda_precomputed_window", - "_replay_each_token", - "_apply_accepted_prefix", - "_delta_rule_step", - "_advance_cursor", - "_fold_prefix", -] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 8bee9f5251..6c19949ef1 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -115,7 +115,7 @@ def _decode_replay( key_slot = req * KH + hv // (HV // KH) sp = State + slot * K * V + kk[:, None] * V + vv[None, :] state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) - # 和半区方案一样：放不下时先折回，再在同一 kernel 里重放新 token。各层只读 cursor。 + # 放不下时先折回，再在同一 kernel 里重放新 token。各层只读 cursor。 if n + (end - start) > L: state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) state = state.to(State.dtype.element_ty).to(tl.float32) @@ -158,8 +158,6 @@ def _decode_replay( start, state, ) - else: - projections, tt = None, None if PRECOMPUTE_PROJECTIONS and KDA: _kda_precomputed_window( @@ -276,8 +274,6 @@ def _decode_replay( K, V, L, - PRECOMPUTE_PROJECTIONS, - BW, KDA, LOWER, BK, @@ -289,11 +285,9 @@ def _decode_replay( kk, ll, n, - projections, slot, start, state, - tt, vv, ) @@ -709,8 +703,6 @@ def _replay_each_token( K: tl.constexpr, V: tl.constexpr, L: tl.constexpr, - PRECOMPUTE_PROJECTIONS: tl.constexpr, - BW: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, BK: tl.constexpr, @@ -722,14 +714,12 @@ def _replay_each_token( kk, ll, n, - projections, slot, start, state, - tt, vv, ): - # KDA、inline 以及不满足 precompute 条件的路径使用逐 token 递推。 + # 逐 token 递推。precompute 走单独的窗口 kernel，不会进这里。 history = ll < n keys = tl.load( Keys + (key_slot * L) * K + ll[:, None] * K + kk[None, :], @@ -777,12 +767,8 @@ def _replay_each_token( keys *= tl.exp(g)[None, :] hk = _history_coeff(keys, k, BK) hq = _history_coeff(keys, q, BK) - 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) + checkpoint_sk = tl.sum(state * (k * tl.exp(total_g))[:, None], axis=0) + checkpoint_sq = tl.sum(state * (q * tl.exp(total_g))[:, None], axis=0) sk = checkpoint_sk + _apply_deltas(hk, deltas, L, BV) sq = checkpoint_sq + _apply_deltas(hq, deltas, L, BV) else: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 0d1af0618b..191a24f32d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -26,11 +26,11 @@ def _fold_prefix( BK: tl.constexpr, BV: tl.constexpr, KDA: tl.constexpr, - STRIDE_LAYER: tl.constexpr = 0, - STRIDE_SLOT: tl.constexpr = 0, - STRIDE_HEAD: tl.constexpr = 0, - STRIDE_K: tl.constexpr = 0, - STRIDE_V: tl.constexpr = 1, + STRIDE_LAYER: tl.constexpr, + STRIDE_SLOT: tl.constexpr, + STRIDE_HEAD: tl.constexpr, + STRIDE_K: tl.constexpr, + STRIDE_V: tl.constexpr, ): # 将 cursor 指向的 accepted history 写回 canonical state。 iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -44,20 +44,16 @@ def _fold_prefix( # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 layer64 = linear_layer_idx.to(tl.int64) head64 = head.to(tl.int64) - stride_layer: tl.constexpr = STRIDE_LAYER if STRIDE_LAYER != 0 else SLOTS * HV * K * V - stride_slot: tl.constexpr = STRIDE_SLOT if STRIDE_SLOT != 0 else HV * K * V - stride_head: tl.constexpr = STRIDE_HEAD if STRIDE_HEAD != 0 else K * V - stride_k: tl.constexpr = STRIDE_K if STRIDE_K != 0 else V slot = (linear_layer_idx * SLOTS + req) * HV + head key_slot = (linear_layer_idx * SLOTS + req) * KH + head // (HV // KH) kk = tl.arange(0, BK) vv = iv * BV + tl.arange(0, BV) sp = ( State - + layer64 * stride_layer - + req * stride_slot - + head64 * stride_head - + kk[:, None].to(tl.int64) * stride_k + + layer64 * STRIDE_LAYER + + req * STRIDE_SLOT + + head64 * STRIDE_HEAD + + kk[:, None].to(tl.int64) * STRIDE_K + vv[None, :].to(tl.int64) * STRIDE_V ) state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) 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 index 28bfdea54d..c1f7b7ffa5 100644 --- 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 @@ -18,13 +18,13 @@ @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): - from lightllm.common.basemodel.triton_kernel.linear_att import replayssm + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay - kernels = ("_decode_replay", "_advance_cursor") + kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} hashes = {name: set() for name in kernels} launches = [] - for name in kernels: - kernel = getattr(replayssm, name) + for name, kernel in kernels.items(): def record(*args, run=kernel.run, name=name, **kwargs): compiled = run(*args, **kwargs) From 664d543853732018bb3bcd35f3a8061e80bd44e3 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 02:07:28 +0000 Subject: [PATCH 42/54] refactor: name replay inputs by the tensors they address The autotune cache key uses the same values under clearer field names, so old entries are left unused. Co-authored-by: Cursor --- .../attention/linear/replay_gdn_decode.py | 56 +++++----- .../attention/linear/replay_kda_decode.py | 56 +++++----- .../linear_att/replayssm_cursor.py | 20 ++-- .../linear_att/replayssm_decode.py | 46 ++++---- .../linear_att/replayssm_delta.py | 6 +- .../linear_att/replayssm_fold.py | 38 +++---- .../triton_kernel/linear_att/ssm_autotune.py | 101 +++++++++--------- .../linear_att/test_ssm_autotune_memory.py | 4 +- 8 files changed, 164 insertions(+), 163 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index cc7b12b248..7cdc0a7d31 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -13,7 +13,7 @@ def replay( - history, + state, linear_layer_idx: int, q: torch.Tensor, k: torch.Tensor, @@ -26,25 +26,25 @@ def replay( cu_seqlens: torch.Tensor | None = None, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 GDN replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + """执行一层 GDN replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。""" - assert not history.kda, "GDN replay step does not use the KDA gate" + assert not state.kda, "GDN replay step does not use the KDA gate" # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 config = ( _replay_config(run_config) if run_config is not None - else resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + else resolve_run_config(state, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) ) - hv, kd, vd = history.ssm_state.shape[-3:] - assert history.num_key_heads in (q.shape[-2], hv) - axis = 1 if cu_seqlens is not None else 0 + hv, kd, vd = state.ssm_state.shape[-3:] + assert state.num_key_heads in (q.shape[-2], hv) + token_axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 - precompute_projections = history.compute_mode == "precompute" and bv >= 16 + precompute_projections = state.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 - bw = triton.next_power_of_2(history.verify_width) - sync_mark = history.tile_sync_mark + padded_verify_width = triton.next_power_of_2(state.verify_width) + sync_mark = state.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -53,39 +53,39 @@ def replay( b, a_log, bias, - history.ssm_state[linear_layer_idx], - history.keys[linear_layer_idx], - history.deltas[linear_layer_idx], - history.gates[linear_layer_idx], - history.raw_keys[linear_layer_idx], - history.raw_values[linear_layer_idx], - history.betas[linear_layer_idx], + state.ssm_state[linear_layer_idx], + state.keys[linear_layer_idx], + state.deltas[linear_layer_idx], + state.gates[linear_layer_idx], + state.raw_keys[linear_layer_idx], + state.raw_values[linear_layer_idx], + state.betas[linear_layer_idx], reqs, - history.cursors, + state.cursors, sync_mark, cu_seqlens, out, - q.stride(axis), - k.stride(axis), - v.stride(axis), + q.stride(token_axis), + k.stride(token_axis), + v.stride(token_axis), a.stride(0), b.stride(0), q.shape[-2], hv, - history.num_key_heads, + state.num_key_heads, kd, vd, - history.capacity, - history.hold_req_id, + state.capacity, + state.hold_req_id, triton.next_power_of_2(kd), bv, cu_seqlens is not None, precompute_projections, - bw, + padded_verify_width, False, - history.gate_lower_bound, - BQK=max(16, 2 * bw), - BR=triton.next_power_of_2(max(16, history.capacity + bw)), + state.gate_lower_bound, + BQK=max(16, 2 * padded_verify_width), + BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 8c2684140c..074ef6fa71 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -13,7 +13,7 @@ def replay( - history, + state, linear_layer_idx: int, q: torch.Tensor, k: torch.Tensor, @@ -26,25 +26,25 @@ def replay( cu_seqlens: torch.Tensor | None = None, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 KDA replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + """执行一层 KDA replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。""" - assert history.kda, "KDA replay step requires the KDA gate" + assert state.kda, "KDA replay step requires the KDA gate" # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 config = ( _replay_config(run_config) if run_config is not None - else resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + else resolve_run_config(state, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) ) - hv, kd, vd = history.ssm_state.shape[-3:] - assert history.num_key_heads in (q.shape[-2], hv) - axis = 1 if cu_seqlens is not None else 0 + hv, kd, vd = state.ssm_state.shape[-3:] + assert state.num_key_heads in (q.shape[-2], hv) + token_axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 - precompute_projections = history.compute_mode == "precompute" and bv >= 16 + precompute_projections = state.compute_mode == "precompute" and bv >= 16 # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 - bw = triton.next_power_of_2(history.verify_width) - sync_mark = history.tile_sync_mark + padded_verify_width = triton.next_power_of_2(state.verify_width) + sync_mark = state.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -53,39 +53,39 @@ def replay( b, a_log, bias, - history.ssm_state[linear_layer_idx], - history.keys[linear_layer_idx], - history.deltas[linear_layer_idx], - history.gates[linear_layer_idx], - history.raw_keys[linear_layer_idx], - history.raw_values[linear_layer_idx], - history.betas[linear_layer_idx], + state.ssm_state[linear_layer_idx], + state.keys[linear_layer_idx], + state.deltas[linear_layer_idx], + state.gates[linear_layer_idx], + state.raw_keys[linear_layer_idx], + state.raw_values[linear_layer_idx], + state.betas[linear_layer_idx], reqs, - history.cursors, + state.cursors, sync_mark, cu_seqlens, out, - q.stride(axis), - k.stride(axis), - v.stride(axis), + q.stride(token_axis), + k.stride(token_axis), + v.stride(token_axis), a.stride(0), b.stride(0), q.shape[-2], hv, - history.num_key_heads, + state.num_key_heads, kd, vd, - history.capacity, - history.hold_req_id, + state.capacity, + state.hold_req_id, triton.next_power_of_2(kd), bv, cu_seqlens is not None, precompute_projections, - bw, + padded_verify_width, True, - history.gate_lower_bound, - BQK=max(16, 2 * bw), - BR=triton.next_power_of_2(max(16, history.capacity + bw)), + state.gate_lower_bound, + BQK=max(16, 2 * padded_verify_width), + BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py index ea3121b64b..9b1ba2fc37 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -21,16 +21,16 @@ def _advance_cursor( BLOCK: tl.constexpr, ): # 先按整窗全接受。放不下时新 token 占 [0, window)，cursor 记为 window；否则加上 window。 - 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) - window = tl.load(cu + i + 1, valid, 0) - tl.load(cu + i, valid, 0) + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(reqs + row, row < N, HOLD) + valid = (row < N) & (req != HOLD) + cursor = tl.load(cursors + req, valid, 0) + window = tl.load(cu + row + 1, valid, 0) - tl.load(cu + row, valid, 0) updated = tl.where(cursor + window > CAPACITY, window, cursor + window) - tl.store(cursors + r, updated, valid) + tl.store(cursors + req, updated, valid) -def advance_cursor(history, reqs, cu_seqlens) -> None: +def advance_cursor(state, reqs, cu_seqlens) -> None: """最后一层返回后，把 cursor 推进到整窗都被接受的位置。 cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。窗口长度为 0 时 cursor 保持不变。HOLD 请求跳过。 @@ -41,12 +41,12 @@ def advance_cursor(history, reqs, cu_seqlens) -> None: assert reqs.is_cuda and cu_seqlens.is_cuda, "advance_cursor requires CUDA inputs" assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( - history.cursors, + state.cursors, reqs, cu_seqlens, reqs.numel(), - history.hold_req_id, - history.capacity, + state.hold_req_id, + state.capacity, 256, num_warps=4, ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 6c19949ef1..f60cac53cd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -47,8 +47,8 @@ def _apply_deltas(coeff, deltas, R: tl.constexpr, BV: tl.constexpr): @triton.jit def _decode_replay( Q, - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -131,7 +131,7 @@ def _decode_replay( projection_precision: tl.constexpr = "tf32x3" qk, projections, tt, tokens, valid, h = _project_qk( Q, - Kp, + StepK, A, Alog, Bias, @@ -161,8 +161,8 @@ def _decode_replay( if PRECOMPUTE_PROJECTIONS and KDA: _kda_precomputed_window( - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -204,8 +204,8 @@ def _decode_replay( ) elif PRECOMPUTE_PROJECTIONS: _gdn_precomputed_window( - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -250,8 +250,8 @@ def _decode_replay( else: _replay_each_token( Q, - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -295,7 +295,7 @@ def _decode_replay( @triton.jit def _project_qk( Q, - Kp, + StepK, A, Alog, Bias, @@ -334,7 +334,7 @@ def _project_qk( 0, ).to(tl.float32) k_rows = tl.load( - Kp + safe_tokens[:, None] * SK + h * K + kk[None, :], + StepK + safe_tokens[:, None] * SK + h * K + kk[None, :], (valid & (tt >= BW))[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) @@ -385,8 +385,8 @@ def _project_qk( @triton.jit def _gdn_precomputed_window( - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -444,7 +444,7 @@ def _gdn_precomputed_window( 0, ).to(tl.float32) raw_keys = tl.load( - Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + StepK + safe_fresh[:, None] * SK + h * K + kk[None, :], fresh[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) @@ -452,7 +452,7 @@ def _gdn_precomputed_window( fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) all_keys = cached_keys + fresh_keys values = tl.load( - Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + StepV + safe_fresh[:, None] * SV + hv * V + vv[None, :], fresh[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) @@ -524,8 +524,8 @@ def _gdn_precomputed_window( @triton.jit def _kda_precomputed_window( - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -581,7 +581,7 @@ def _kda_precomputed_window( 0, ).to(tl.float32) raw_keys = tl.load( - Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + StepK + safe_fresh[:, None] * SK + h * K + kk[None, :], fresh[:, None] & (kk[None, :] < K), 0, ).to(tl.float32) @@ -589,7 +589,7 @@ def _kda_precomputed_window( fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) all_keys = tl.where(old[:, None], cached_keys, 0) + fresh_keys values = tl.load( - Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + StepV + safe_fresh[:, None] * SV + hv * V + vv[None, :], fresh[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) @@ -679,8 +679,8 @@ def _kda_precomputed_window( @triton.jit def _replay_each_token( Q, - Kp, - Vp, + StepK, + StepV, A, B, Alog, @@ -749,8 +749,8 @@ def _replay_each_token( 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) + raw_k = tl.load(StepK + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(StepV + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) if KDA: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py index 468246bf18..e832c6f3f3 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -32,15 +32,15 @@ def _apply_accepted_prefix( slot, kk, vv, - n, + prefix_len, K: tl.constexpr, V: tl.constexpr, L: tl.constexpr, KDA: tl.constexpr, ): - # 按 history 槽位顺序把已接受前缀折进 state。j >= n 的槽不读、不更新。 + # 按 history 槽位顺序把已接受前缀折进 state。j >= prefix_len 的槽不读、不更新。 for j in range(L): - active = j < n + active = j < prefix_len raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) if KDA: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py index 191a24f32d..42aad6edc2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -37,8 +37,8 @@ def _fold_prefix( req = tl.load(Reqs + row).to(tl.int64) if req == HOLD: return - n = tl.load(Cursors + req) - if n == 0: + prefix_len = tl.load(Cursors + req) + if prefix_len == 0: return linear_layer_idx, head = lh // HV, lh % HV # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 @@ -58,13 +58,15 @@ def _fold_prefix( ) state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 - state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) + state = _apply_accepted_prefix( + state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, prefix_len, K, V, L, KDA + ) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) # cursor 不在这里清。和折回并行的其他线程块还要读原来的长度，清零放到 kernel 结束之后。 -def fold_prefix(history, reqs, run_config=None) -> None: - """把 history 上 cursor 覆盖的已接受前缀折进 canonical SSM。 +def fold_prefix(state, reqs, run_config=None) -> None: + """把 cursor 覆盖的已接受前缀折进 canonical SSM。 没传入 run_config 时用 BV=32、4 个 warp、3 个 stage。 """ @@ -74,27 +76,27 @@ def fold_prefix(history, reqs, run_config=None) -> None: bv = run_config["BV"] num_warps = run_config["num_warps"] num_stages = run_config["num_stages"] - state = history.ssm_state - layers, slots, hv, k, v = state.shape - layer_stride, slot_stride, head_stride, key_stride, value_stride = state.stride() + ssm = state.ssm_state + layers, slots, hv, k, v = ssm.shape + layer_stride, slot_stride, head_stride, key_stride, value_stride = ssm.stride() _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( - state, - history.raw_keys, - history.raw_values, - history.gates, - history.betas, - history.cursors, + ssm, + state.raw_keys, + state.raw_values, + state.gates, + state.betas, + state.cursors, reqs, slots, hv, - history.num_key_heads, + state.num_key_heads, k, v, - history.capacity, - history.hold_req_id, + state.capacity, + state.hold_req_id, triton.next_power_of_2(k), bv, - history.kda, + state.kda, STRIDE_LAYER=layer_stride, STRIDE_SLOT=slot_stride, STRIDE_HEAD=head_stride, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index ce761a3483..38a0a3f40a 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -16,19 +16,19 @@ _DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} -def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: +def resolve_run_config(state, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: """返回当前 batch 的固定或已调优 SSM kernel 配置。""" - if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: - return _replay_config(ssm_backend.run_config) - default = _replay_config(ssm_backend.run_config) - key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) - if key in ssm_backend.batch_configs: - return ssm_backend.batch_configs[key] + if state._config_is_fixed and state.run_config is not None: + return _replay_config(state.run_config) + default = _replay_config(state.run_config) + key = (tuple(static_key(state, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) + if key in state.batch_configs: + return state.batch_configs[key] if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): return default - config = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) + config = _replay_config(select_config(state, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 - ssm_backend.batch_configs[key] = config + state.batch_configs[key] = config logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") return config @@ -42,60 +42,58 @@ def get_configs() -> list[dict[str, int]]: ] -def static_key(ssm_backend, mode, q, cu_seqlens) -> dict[str, object]: +def static_key(state, mode, q, cu_seqlens) -> dict[str, object]: """构造区分状态布局、数据类型和投影模式的静态调优键。""" - _, _, hv, kd, vd = ssm_backend.ssm_state.shape + _, _, hv, kd, vd = state.ssm_state.shape # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 key = { "mode": mode, - "v": 7, - "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), - "dtype": str((q.dtype, ssm_backend.ssm_state.dtype)), - "history": ssm_backend.capacity, - "gate": int(ssm_backend.kda), + "v": 8, + "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), + "dtype": str((q.dtype, state.ssm_state.dtype)), + "capacity": state.capacity, + "kda": int(state.kda), "varlen": cu_seqlens is not None, } - if ssm_backend.kda: - key["lower"] = ssm_backend.gate_lower_bound - key["history_dtype"] = str(ssm_backend.keys.dtype) - key["projection"] = ( - ("kda_mma" if ssm_backend.kda else "gdn_mma") if ssm_backend.compute_mode == "precompute" else "inline" - ) + if state.kda: + key["lower"] = state.gate_lower_bound + key["recon_dtype"] = str(state.keys.dtype) + key["projection"] = ("kda_mma" if state.kda else "gdn_mma") if state.compute_mode == "precompute" else "inline" return key def run_key(q, cu_seqlens) -> tuple[int, int]: - """返回调优使用的 token 数和请求数键。""" + """返回调优键：token 数，以及请求数。没有 cu_seqlens 时第二项也是 token 数，因为一行一个请求。""" 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 rebuild_inputs( - ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None + state, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None ) -> tuple[tuple[object, ...], dict[str, object]]: """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" - tokens, sequences = run_key(q, cu_seqlens) - batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens - _, _, hv, kd, vd = ssm_backend.ssm_state.shape + tokens, run_count = run_key(q, cu_seqlens) + batch = min(run_count, triton.cdiv(tokens, state.verify_width)) if cu_seqlens is not None else tokens + _, _, hv, kd, vd = state.ssm_state.shape from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig # 计时只需要一层、本 batch 的槽位。正式执行仍用调用方原来的 state。 - scratch = type(ssm_backend)( + scratch = type(state)( ConvSSMStateConfig( logical_layer_ids=(0,), - global_k_heads=ssm_backend.num_key_heads, + global_k_heads=state.num_key_heads, global_v_heads=hv, key_dim=kd, value_dim=vd, conv_width=1, conv_dtype=q.dtype, - ssm_dtype=ssm_backend.ssm_state.dtype, + ssm_dtype=state.ssm_state.dtype, ), batch + 1, - mtp_step=ssm_backend.verify_width - 1, + mtp_step=state.verify_width - 1, device=q.device, - ssm_replay_config={"cache_len": ssm_backend.capacity, "compute_mode": ssm_backend.compute_mode}, - gate_lower_bound=ssm_backend.gate_lower_bound, + ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, + gate_lower_bound=state.gate_lower_bound, ) scratch.ssm_state.normal_(0, 0.02) # benchmark 回调不能再次触发配置选择。 @@ -106,12 +104,12 @@ def rebuild_inputs( reqs = torch.cat( ( torch.arange(batch, dtype=torch.int32, device=q.device), - torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), + torch.full((run_count - 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) + cu = (torch.arange(run_count + 1, dtype=torch.int32, device=q.device) * state.verify_width).clamp_max(tokens) return (), dict( - ssm_backend=scratch, + state=scratch, mode=mode, q=inputs[0], k=inputs[1], @@ -125,19 +123,19 @@ def rebuild_inputs( ) -def replay_layer(history, *args, **kwargs): - """按 history 的具体类型调用对应的 replay decode。""" +def replay_layer(state, *args, **kwargs): + """按 state 的具体类型调用对应的 replay decode。""" from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - if isinstance(history, ReplayKDAState): + if isinstance(state, ReplayKDAState): from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay - elif isinstance(history, ReplayGDNState): + elif isinstance(state, ReplayGDNState): from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay else: - raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(history).__name__}") - return replay(history, *args, **kwargs) + raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(state).__name__}") + return replay(state, *args, **kwargs) @autotune( @@ -151,23 +149,24 @@ def replay_layer(history, *args, **kwargs): 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 + state, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None ) -> dict[str, int] | None: """搜索 update cycle 布局；计时只使用重建的临时输入。""" if workload is None: return run_config - ssm_backend.run_config = run_config + state.run_config = run_config (reqs,) = workload - ssm_backend.cursors.zero_() + state.cursors.zero_() # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。调优窗口全部接受。 tokens = q.shape[1 if cu_seqlens is not None else 0] - rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) + rounds = max(1, state.capacity // min(state.verify_width, tokens)) for _ in range(rounds): - replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) - adopt_cu = cu_seqlens - if adopt_cu is None: - adopt_cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - advance_cursor(ssm_backend, reqs, adopt_cu) + replay_layer(state, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + # 宽度为 1 时 replay 不用 cu_seqlens，cursor 仍要一份每行一个 token 的前缀和。 + cursor_cu = cu_seqlens + if cursor_cu is None: + cursor_cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + advance_cursor(state, reqs, cursor_cu) return run_config 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 index 8f3f8ec976..85e83fa150 100644 --- 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 @@ -18,7 +18,7 @@ def test_rebuild_uses_one_layer_scratch(width): a_log = torch.zeros(2, device="cuda") cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) - scratch = kwargs["ssm_backend"] + scratch = kwargs["state"] assert scratch is not ssm_backend assert scratch.ssm_state.shape[0] == 1 assert scratch.capacity == ssm_backend.capacity @@ -45,7 +45,7 @@ def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, w 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 + state=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)} From 88b886e21ae889d828bc1588adcc8e5b7ef29bb7 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 02:45:26 +0000 Subject: [PATCH 43/54] refactor: require one packed layout for replay decode Decode reads every request window from cu_seqlens, and q, k, v must already be [tokens, heads, dim]. Autotune times that same layout, so old cache entries are left unused. Co-authored-by: Cursor --- .../basemodel/attention/linear/replay_gdn.py | 18 +++-- .../attention/linear/replay_gdn_decode.py | 18 +++-- .../basemodel/attention/linear/replay_kda.py | 19 +++--- .../attention/linear/replay_kda_decode.py | 18 +++-- .../linear_att/replayssm_decode.py | 6 +- .../triton_kernel/linear_att/ssm_autotune.py | 24 ++++--- test/benchmark/kernels/benchmark_kda_ssm.py | 25 ++++--- .../linear_att/test_acceptance_fusion.py | 4 +- .../linear_att/test_replayssm.py | 66 ++++++++++--------- .../linear_att/test_replayssm_execution.py | 26 ++++---- .../linear_att/test_ssm_autotune_memory.py | 8 +-- .../common/test_linear_state_contract.py | 6 +- 12 files changed, 126 insertions(+), 112 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 3892648934..ead3cecd42 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -201,7 +201,7 @@ def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None self.b_conv_buffer_idx = infer.b_req_idx self.b1_mtp_cu_q_seq_len = None self.b_num_accepted_tokens = None - # 每行一个 token。replay kernel 仍用这个布局，cursor 需要现成的前缀和。 + # 每行一个 token。replay 和 cursor 共用这份前缀和。 self.b_cu_seqlens = torch.arange( infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device ) @@ -250,6 +250,9 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat backend.tp_num_v_heads, backend.head_v_dim, ) + query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) + key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) + value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) output, _ = replay( state, state.storage_layer(layer_num), @@ -261,7 +264,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, + self.b_cu_seqlens, ) return output, z @@ -278,7 +281,14 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step num_accepted_tokens=self.b_num_accepted_tokens, query_start_loc=self.b1_mtp_cu_q_seq_len, ) - query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) output, _ = replay( state, state.storage_layer(layer_num), @@ -290,7 +300,7 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, + self.b_cu_seqlens, ) return output diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 7cdc0a7d31..d7a8fcd5ab 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -23,12 +23,18 @@ def replay( a_log: torch.Tensor, bias: torch.Tensor, reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 GDN replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。""" + """执行一层 GDN replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 + + q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 + """ assert not state.kda, "GDN replay step does not use the KDA gate" + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 config = ( _replay_config(run_config) @@ -37,7 +43,6 @@ def replay( ) hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) - token_axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 @@ -65,9 +70,9 @@ def replay( sync_mark, cu_seqlens, out, - q.stride(token_axis), - k.stride(token_axis), - v.stride(token_axis), + q.stride(0), + k.stride(0), + v.stride(0), a.stride(0), b.stride(0), q.shape[-2], @@ -79,7 +84,6 @@ def replay( state.hold_req_id, triton.next_power_of_2(kd), bv, - cu_seqlens is not None, precompute_projections, padded_verify_width, False, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index 0354cfe5fc..cc6895efb5 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -177,24 +177,23 @@ def decode_att( **conv_kwargs, ) q, k, v = backend.split_qkv(mixed_qkv) - shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) - q = q.view(*shape, backend.tp_num_heads, backend.head_dim) - k = k.view(*shape, backend.tp_num_heads, backend.head_dim) - v = v.view(*shape, backend.tp_num_heads, backend.head_dim) - raw_gate = params["raw_gate"].view(*shape, backend.tp_hidden_size) - raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) + q = q.view(-1, backend.tp_num_heads, backend.head_dim) + k = k.view(-1, backend.tp_num_heads, backend.head_dim) + v = v.view(-1, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) + raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) output, _ = replay( state, state.storage_layer(layer_num), q, k, v, - raw_gate.reshape(-1, backend.tp_hidden_size), - raw_beta.reshape(-1, backend.tp_num_heads), + raw_gate, + raw_beta, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, + self.b_cu_seqlens, ) # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) @@ -206,7 +205,7 @@ def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: self.b_conv_buffer_idx = infer.b_req_idx self.b1_mtp_cu_q_seq_len = None self.b_num_accepted_tokens = None - # 每行一个 token。replay kernel 仍用这个布局，cursor 需要现成的前缀和。 + # 每行一个 token。replay 和 cursor 共用这份前缀和。 self.b_cu_seqlens = torch.arange( infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device ) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 074ef6fa71..d76007a73b 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -23,12 +23,18 @@ def replay( a_log: torch.Tensor, bias: torch.Tensor, reqs: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor, run_config: dict[str, int] | None = None, ) -> tuple[torch.Tensor, dict[str, int]]: - """执行一层 KDA replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。""" + """执行一层 KDA replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 + + q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 + """ assert state.kda, "KDA replay step requires the KDA gate" + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 config = ( _replay_config(run_config) @@ -37,7 +43,6 @@ def replay( ) hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) - token_axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 @@ -65,9 +70,9 @@ def replay( sync_mark, cu_seqlens, out, - q.stride(token_axis), - k.stride(token_axis), - v.stride(token_axis), + q.stride(0), + k.stride(0), + v.stride(0), a.stride(0), b.stride(0), q.shape[-2], @@ -79,7 +84,6 @@ def replay( state.hold_req_id, triton.next_power_of_2(kd), bv, - cu_seqlens is not None, precompute_projections, padded_verify_width, True, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index f60cac53cd..8f16764d54 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -79,7 +79,6 @@ def _decode_replay( HOLD: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, - VARLEN: tl.constexpr, PRECOMPUTE_PROJECTIONS: tl.constexpr, BW: tl.constexpr, KDA: tl.constexpr, @@ -95,10 +94,7 @@ def _decode_replay( """ 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 + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) vv = iv * BV + tl.arange(0, BV) if req == HOLD: for t in range(start, end): diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 38a0a3f40a..ab29f3ddd4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -48,12 +48,11 @@ def static_key(state, mode, q, cu_seqlens) -> dict[str, object]: # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 key = { "mode": mode, - "v": 8, + "v": 10, "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), "dtype": str((q.dtype, state.ssm_state.dtype)), "capacity": state.capacity, "kda": int(state.kda), - "varlen": cu_seqlens is not None, } if state.kda: key["lower"] = state.gate_lower_bound @@ -63,8 +62,8 @@ def static_key(state, mode, q, cu_seqlens) -> dict[str, object]: def run_key(q, cu_seqlens) -> tuple[int, int]: - """返回调优键：token 数，以及请求数。没有 cu_seqlens 时第二项也是 token 数，因为一行一个请求。""" - tokens = q.shape[1 if cu_seqlens is not None else 0] + """返回调优键：token 数，以及请求数。q 是 [tokens, heads, dim]。没传前缀和时一行一个请求。""" + tokens = q.shape[0] return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens @@ -73,7 +72,8 @@ def rebuild_inputs( ) -> tuple[tuple[object, ...], dict[str, object]]: """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" tokens, run_count = run_key(q, cu_seqlens) - batch = min(run_count, triton.cdiv(tokens, state.verify_width)) if cu_seqlens is not None else tokens + # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) _, _, hv, kd, vd = state.ssm_state.shape from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig @@ -107,7 +107,9 @@ def rebuild_inputs( torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), ) ) - cu = (torch.arange(run_count + 1, dtype=torch.int32, device=q.device) * state.verify_width).clamp_max(tokens) + # 计时用均匀窗口，不读调用方的前缀和。宽度为 1 时 verify_width 就是 1。 + cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) + cu = (cu * state.verify_width).clamp_max(tokens) return (), dict( state=scratch, mode=mode, @@ -118,7 +120,7 @@ def rebuild_inputs( b=inputs[4], a_log=inputs[5], bias=inputs[6], - cu_seqlens=cu if cu_seqlens is not None else None, + cu_seqlens=cu, workload=(reqs,), ) @@ -158,15 +160,11 @@ def select_config( (reqs,) = workload state.cursors.zero_() # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。调优窗口全部接受。 - tokens = q.shape[1 if cu_seqlens is not None else 0] + tokens = run_key(q, cu_seqlens)[0] rounds = max(1, state.capacity // min(state.verify_width, tokens)) for _ in range(rounds): replay_layer(state, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) - # 宽度为 1 时 replay 不用 cu_seqlens，cursor 仍要一份每行一个 token 的前缀和。 - cursor_cu = cu_seqlens - if cursor_cu is None: - cursor_cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - advance_cursor(state, reqs, cursor_cu) + advance_cursor(state, reqs, cu_seqlens) return run_config diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 706807de5a..39a3d1f388 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -23,15 +23,17 @@ 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)] + q, k, v = [x.view(batch * width, heads, dim) for x in packed.split(heads * dim, -1)] + native_shape = (1, batch * width) if width > 1 else (batch, 1) gate = torch.randn(batch * width, heads * dim, dtype=q.dtype, device="cuda") - 3 beta = torch.randn(batch * width, heads, dtype=q.dtype, device="cuda") alog = torch.zeros(heads, device="cuda") bias = torch.zeros(heads * dim, device="cuda") reqs = torch.arange(batch, dtype=torch.int32, device="cuda") - cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width if width > 1 else None + # 原生 decode 在宽度为 1 时不收前缀和。replay 始终要，一行一个 token 时就是 0, 1, 2, ...。 + replay_cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width + cu = replay_cu if width > 1 else None counts = reqs % width + 1 if acceptance == "partial" else torch.full_like(reqs, width) accepted = torch.cat((counts - 1, counts.new_zeros(1))) snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") @@ -59,11 +61,11 @@ def cycle(): 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), + q.view(*native_shape, heads, dim), + k.view(*native_shape, heads, dim), + v.view(*native_shape, heads, dim), + gate.view(*native_shape, -1), + beta.view(*native_shape, heads), alog, bias, state[layer], @@ -72,12 +74,9 @@ def cycle(): num_accepted_tokens=counts if cu is not None else None, ) else: - replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, cu) + replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, replay_cu) if ssm_backend is not None: - if cu is None: - accept_replay_window(ssm_backend, reqs) - else: - accept_replay_window(ssm_backend, reqs, accepted, cu) + accept_replay_window(ssm_backend, reqs, accepted, replay_cu) if isinstance(ssm_backend, ReplaySSMExecutor): ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) ssm_backend.fold_accepted_prefix(reqs) 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 index f6cac54b6a..f228eb7133 100644 --- 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 @@ -188,9 +188,9 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput 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") + q = torch.randn(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") + v = torch.randn(12, 4, 32, dtype=torch.bfloat16, device="cuda") a = torch.full((12, 4), -3.0, dtype=torch.bfloat16, device="cuda") b = torch.randn_like(a) alog = torch.zeros(4, device="cuda") 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 index 53f38e15d1..894bcd5276 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -64,13 +64,13 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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 + q = torch.zeros(4, 1, 128, dtype=torch.float16, device="cuda") + q[:2, 0, 0] = 1 k = torch.zeros_like(q) - k[0, 0, 0, 0] = -0.25 - k[0, 1, 0, 0] = 1 + k[0, 0, 0] = -0.25 + k[1, 0, 0] = 1 v = torch.zeros_like(q) - v[0, :2, 0, 0] = 60000 + v[: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") @@ -86,12 +86,12 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state 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 + q[token], k[token], v[token], a[token], b[token], alog, bias, expected_state, kda=kda ) expected.append(output) assert torch.isfinite(expected_state).all() assert torch.isfinite(actual).all() - torch.testing.assert_close(actual[0, :2].float(), torch.stack(expected), atol=4, rtol=1e-3) + torch.testing.assert_close(actual[:2].float(), torch.stack(expected), atol=4, rtol=1e-3) accept_replay_window(ssm_backend, reqs, accepted, cu) # The correction is larger than FP16 can represent even though the # recurrent state, inputs, and attention outputs are all finite. @@ -175,9 +175,9 @@ def test_replay_acceptance_flush_and_graph( 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) + q = torch.randn(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) + v = torch.randn(tokens, hv, vdim, device="cuda", dtype=activation_dtype) a = torch.randn(tokens, hv * kdim if kda else hv, device="cuda", dtype=activation_dtype) + gate_shift b = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) alog = torch.randn(hv, device="cuda") * 0.1 @@ -220,9 +220,9 @@ def step(): for j in range(length): t = start + j out, cur = reference( - q[0, t], - k[0, t], - v[0, t], + q[t], + k[t], + v[t], a[t], b[t], alog, @@ -231,7 +231,7 @@ def step(): kda=kda, gate_lower_bound=gate_lower_bound, ) - torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) if j + 1 == counts[seq]: expected[layer, req].copy_(cur) if iteration == 2 * capacity + 1: @@ -330,9 +330,9 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d 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) + q = torch.randn((tokens, 2, 128), device="cuda", dtype=dtype) k = torch.randn_like(q) - v = torch.randn((1, tokens, 4, 128), device="cuda", dtype=dtype) + v = torch.randn((tokens, 4, 128), device="cuda", dtype=dtype) a = torch.randn((tokens, 4), device="cuda", dtype=dtype) + decay b = torch.randn_like(a) alog = torch.zeros(4, device="cuda") @@ -391,9 +391,10 @@ def test_bf16_commit_rounds_before_the_next_token(): alog = torch.zeros(1, device="cuda") bias = torch.zeros_like(alog) outputs = [] + cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: assert ssm_backend.cursors[0].item() == 0 - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs)[0]) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0]) torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) @@ -422,13 +423,13 @@ def check(capacity, width, steps): ) for mode in ("inline", "precompute") ] - cu = None if width == 1 else torch.tensor([0, width], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, width], dtype=torch.int32, device="cuda") accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") for step in range(steps): torch.manual_seed(100 + step) - q = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + q = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) - v = torch.randn((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 b = torch.randn_like(a) outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] for backend in backends] @@ -460,6 +461,7 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) assert backend.gates.shape[-1] == (dim if kda else capacity) reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") alog = torch.zeros(hv, device="cuda") bias = torch.zeros(hv * dim if kda else hv, device="cuda") for step in range(capacity + 1): @@ -474,8 +476,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): assert int(backend.cursors[0]) == capacity outs, twin_outs = [], [] for layer in range(2): - outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs)[0]) - twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs)[0]) + outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0]) + twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu)[0]) if step == capacity: assert int(backend.cursors[0]) == capacity for actual, expected in zip(outs, twin_outs): @@ -498,9 +500,9 @@ def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): backend.cursors[0] = 3 reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) - q = torch.randn(1, 1, 2, 32, device="cuda") + q = torch.randn(1, 2, 32, device="cuda") k = torch.randn_like(q) - v = torch.randn(1, 1, 2, 32, device="cuda") + v = torch.randn(1, 2, 32, device="cuda") a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") b = torch.zeros(1, 2, device="cuda") bias = torch.zeros(2 * 32 if kda else 2, device="cuda") @@ -532,9 +534,9 @@ def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_ 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") + q = torch.randn(8, 2, 128, dtype=dtype, device="cuda") k = torch.randn_like(q) - v = torch.randn(1, 8, 4, 128, dtype=dtype, device="cuda") + v = torch.randn(8, 4, 128, dtype=dtype, device="cuda") a = torch.full((8, 4 * 128 if kda else 4), -3.0, dtype=dtype, device="cuda") b = torch.randn(8, 4, dtype=dtype, device="cuda") alog = torch.zeros(4, device="cuda") @@ -590,8 +592,8 @@ def step(): 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, bias, current, kda=kda) - torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + out, current = reference(q[t], k[t], v[t], a[t], b[t], alog, bias, current, kda=kda) + torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) if j + 1 == count: expected[layer, req].copy_(current) pending[req] += count @@ -773,9 +775,9 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa actual = replay_layer( backend, 0, - q.reshape(1, -1, 1, 32), - k.reshape(1, -1, 1, 32), - v.reshape(1, -1, 1, 32), + q.reshape(-1, 1, 32), + k.reshape(-1, 1, 32), + v.reshape(-1, 1, 32), a.flatten(0, 1), b.flatten(0, 1), alog, @@ -821,9 +823,9 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): bias = torch.zeros(hv * dim if kda else hv, device="cuda") cu = torch.tensor([0, 2], device="cuda", dtype=torch.int32) for step in range(3): - q = torch.randn(1, 2, kh, dim, device="cuda", dtype=dtype) + q = torch.randn(2, kh, dim, device="cuda", dtype=dtype) k = torch.randn_like(q) - v = torch.randn(1, 2, hv, dim, device="cuda", dtype=dtype) + v = torch.randn(2, hv, dim, device="cuda", dtype=dtype) a = torch.full((2, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) b = torch.randn(2, hv, device="cuda", dtype=dtype) replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) 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 index c1f7b7ffa5..e4693181e0 100644 --- 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 @@ -49,8 +49,9 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() - out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] - accept_replay_window(backend, ids) + cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) + out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids, cu)[0] + accept_replay_window(backend, ids, cu_seqlens=cu) return out graphs = {} @@ -115,9 +116,9 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle 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) + q = torch.randn(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) + v = torch.randn(12, 4, 128, device="cuda", dtype=torch.bfloat16) a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) alog = torch.zeros(4, device="cuda") @@ -259,8 +260,10 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, alog = torch.zeros(4, device="cuda") bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + def forward(ssm_backend): - return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] + return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] def warm(ssm_backend): out = forward(ssm_backend) @@ -334,9 +337,9 @@ def test_cross_layer_fold_with_state_offset_above_int32(kda): ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda, run_config=config) for state in [dense, sparse] ] - q = torch.randn(1, 8, 4, 128, device="cuda", dtype=torch.bfloat16) + q = torch.randn(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) + v = torch.randn(8, 12, 128, device="cuda", dtype=torch.bfloat16) a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) alog = torch.zeros(12, device="cuda") @@ -385,9 +388,9 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): 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) + q = torch.randn(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) + v = torch.randn(batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) alog = torch.zeros(4, device="cuda") @@ -446,6 +449,7 @@ def test_full_prefix_does_not_address_past_capacity(kda): run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) q = torch.randn(2, 1, 32, device="cuda") * 0.2 k, v = torch.randn_like(q), torch.randn_like(q) a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") @@ -456,13 +460,13 @@ def test_full_prefix_does_not_address_past_capacity(kda): for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): getattr(backend, name).zero_() backend.cursors[reqs.long()] = capacity + 2 - replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs, cu) backend.fold_accepted_prefix(reqs) torch.cuda.synchronize() assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) assert torch.isfinite(backend.ssm_state).all() for _ in range(12): - replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs) + replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs, cu) accept_replay_window(backend, reqs) overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] if overflow.numel(): 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 index 85e83fa150..b6c0aca7b0 100644 --- 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 @@ -13,10 +13,10 @@ def test_rebuild_uses_one_layer_scratch(width): state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) tokens = 3 * width - q = torch.ones((1, tokens, 2, 32) if width > 1 else (tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") a_log = torch.zeros(2, device="cuda") - cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) scratch = kwargs["state"] assert scratch is not ssm_backend @@ -41,9 +41,7 @@ def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, w 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) + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width inputs = dict( state=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu ) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index f27909f889..cf89b94480 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -140,8 +140,8 @@ def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, modul ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, ) heads = owner.num_key_heads - q = torch.zeros(1, 1, heads, linear_config.key_dim) - v = torch.zeros(1, 1, owner.ssm_state.shape[-3], linear_config.value_dim) + q = torch.zeros(1, heads, linear_config.key_dim) + v = torch.zeros(1, owner.ssm_state.shape[-3], linear_config.value_dim) gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] decode.replay( owner, @@ -154,7 +154,7 @@ def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, modul torch.zeros(v.shape[-3]), torch.zeros(gate), torch.zeros(1, dtype=torch.int32), - None, + torch.arange(2, dtype=torch.int32), run_config={"BV": bv, "num_warps": 1, "num_stages": 1}, ) assert captured.args[33] is precompute From 7bd612f2fadeb38b98ae6e7fda3ca569e7664f2a Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 03:19:51 +0000 Subject: [PATCH 44/54] refactor: tune GDN and KDA replay on their own kernels The shared tuner only forwarded both algorithms through one entry. Each replay kernel now owns its config search and returns only the attention output. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 1 - .../bounded_att/kda/replay_kda_state.py | 1 - .../basemodel/attention/linear/replay_gdn.py | 4 +- .../attention/linear/replay_gdn_decode.py | 136 +++++++++++-- .../basemodel/attention/linear/replay_kda.py | 2 +- .../attention/linear/replay_kda_decode.py | 137 ++++++++++++-- .../triton_kernel/linear_att/ssm_autotune.py | 178 ------------------ test/benchmark/kernels/benchmark_kda_ssm.py | 4 +- .../linear_att/replay_history.py | 11 +- .../linear_att/test_acceptance_fusion.py | 4 +- .../linear_att/test_replayssm.py | 44 +++-- .../linear_att/test_replayssm_execution.py | 18 +- .../linear_att/test_ssm_autotune_memory.py | 16 +- .../common/test_linear_state_contract.py | 12 -- 14 files changed, 293 insertions(+), 275 deletions(-) delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 8508123d50..f972d4b884 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -96,7 +96,6 @@ def __init__( self.compute_mode = compute_mode self.run_config = None self._config_is_fixed = False - self.batch_configs = {} layers, slots, hv, k, v = self.ssm_state.shape self.num_key_heads = config.num_k_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 983eccb5c8..133d6063e8 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -96,7 +96,6 @@ def __init__( self.compute_mode = compute_mode self.run_config = None self._config_is_fixed = False - self.batch_configs = {} layers, slots, hv, k, v = self.ssm_state.shape self.num_key_heads = config.num_k_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index ead3cecd42..5b96d21596 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -253,7 +253,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) - output, _ = replay( + output = replay( state, state.storage_layer(layer_num), query, @@ -289,7 +289,7 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) - output, _ = replay( + output = replay( state, state.storage_layer(layer_num), query, diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index d7a8fcd5ab..fc345fe34e 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -1,17 +1,119 @@ -"""GDN replay decode 的 kernel 辅助。 +"""GDN replay decode。 -History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置，并发起一层重放。 +History 张量和 cursor 在 ReplayGDNState。这里调优并执行一层重放。 """ from __future__ import annotations +import ast + import torch import triton +from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def get_configs() -> list[dict[str, int]]: + """返回 GDN replay 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def get_static_key(state, q) -> dict[str, object]: + """区分头维度、状态精度、history 容量和投影模式。""" + _, _, hv, kd, vd = state.ssm_state.shape + return { + "v": 1, + "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), + "dtype": str((q.dtype, state.ssm_state.dtype)), + "capacity": state.capacity, + "recon_dtype": str(state.keys.dtype), + "projection": "gdn_mma" if state.compute_mode == "precompute" else "inline", + } + + +def get_run_key(q, cu_seqlens) -> tuple[int, int]: + """调优键是 token 数和请求数。q 是 [tokens, heads, dim]。""" + return q.shape[0], cu_seqlens.numel() - 1 + + +def _check_replay_inputs(q, k, v, reqs, cu_seqlens) -> None: + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" +def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens, run_config=None): + """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" + del linear_layer_idx, reqs, run_config + tokens, run_count = get_run_key(q, cu_seqlens) + # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) + _, _, hv, kd, vd = state.ssm_state.shape + scratch = type(state)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=state.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=state.ssm_state.dtype, + ), + batch + 1, + mtp_step=state.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, + gate_lower_bound=state.gate_lower_bound, + ) + scratch.ssm_state.normal_(0, 0.02) + copied = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + bench_reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + # 计时用均匀窗口。宽度为 1 时 verify_width 就是 1。 + cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) + cu = (cu * state.verify_width).clamp_max(tokens) + return (), dict( + state=scratch, + linear_layer_idx=0, + q=copied[0], + k=copied[1], + v=copied[2], + a=copied[3], + b=copied[4], + a_log=copied[5], + bias=copied[6], + reqs=bench_reqs, + cu_seqlens=cu, + ) + + +def _run_key_distance(left: str, right: str) -> int: + return sum(abs(x - y) for x, y in zip(ast.literal_eval(left), ast.literal_eval(right))) + + +@autotune( + kernel_name="replay_gdn_decode:v1", + configs_gen_func=get_configs, + static_key_func=get_static_key, + run_key_func=get_run_key, + run_key_distance_func=_run_key_distance, + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) def replay( state, linear_layer_idx: int, @@ -25,31 +127,25 @@ def replay( reqs: torch.Tensor, cu_seqlens: torch.Tensor, run_config: dict[str, int] | None = None, -) -> tuple[torch.Tensor, dict[str, int]]: +) -> torch.Tensor: """执行一层 GDN replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 """ assert not state.kda, "GDN replay step does not use the KDA gate" - assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" - assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" - assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" - # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 - config = ( - _replay_config(run_config) - if run_config is not None - else resolve_run_config(state, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - ) + _check_replay_inputs(q, k, v, reqs, cu_seqlens) + if state._config_is_fixed and state.run_config is not None: + run_config = state.run_config + elif run_config is None: + run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) out = torch.empty_like(v) - bv = config["BV"] + bv = run_config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 precompute_projections = state.compute_mode == "precompute" and bv >= 16 - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 padded_verify_width = triton.next_power_of_2(state.verify_width) - sync_mark = state.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -67,7 +163,7 @@ def replay( state.betas[linear_layer_idx], reqs, state.cursors, - sync_mark, + state.tile_sync_mark, cu_seqlens, out, q.stride(0), @@ -91,7 +187,7 @@ def replay( BQK=max(16, 2 * padded_verify_width), BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), + num_warps=run_config["num_warps"], + num_stages=run_config["num_stages"], ) - return out, config + return out diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index cc6895efb5..e5e8d4df85 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -182,7 +182,7 @@ def decode_att( v = v.view(-1, backend.tp_num_heads, backend.head_dim) raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) - output, _ = replay( + output = replay( state, state.storage_layer(layer_num), q, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index d76007a73b..50dc383f6f 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -1,17 +1,120 @@ -"""KDA replay decode 的 kernel 辅助。 +"""KDA replay decode。 -History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置，并发起一层重放。 +History 张量和 cursor 在 ReplayKDAState。这里调优并执行一层重放。 """ from __future__ import annotations +import ast + import torch import triton +from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import _replay_config, resolve_run_config +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def get_configs() -> list[dict[str, int]]: + """返回 KDA replay 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def get_static_key(state, q) -> dict[str, object]: + """区分头维度、状态精度、history 容量、gate 下界和投影模式。""" + _, _, hv, kd, vd = state.ssm_state.shape + return { + "v": 1, + "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), + "dtype": str((q.dtype, state.ssm_state.dtype)), + "capacity": state.capacity, + "lower": state.gate_lower_bound, + "recon_dtype": str(state.keys.dtype), + "projection": "kda_mma" if state.compute_mode == "precompute" else "inline", + } + + +def get_run_key(q, cu_seqlens) -> tuple[int, int]: + """调优键是 token 数和请求数。q 是 [tokens, heads, dim]。""" + return q.shape[0], cu_seqlens.numel() - 1 + + +def _check_replay_inputs(q, k, v, reqs, cu_seqlens) -> None: + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" +def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens, run_config=None): + """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" + del linear_layer_idx, reqs, run_config + tokens, run_count = get_run_key(q, cu_seqlens) + # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) + _, _, hv, kd, vd = state.ssm_state.shape + scratch = type(state)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=state.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=state.ssm_state.dtype, + ), + batch + 1, + mtp_step=state.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, + gate_lower_bound=state.gate_lower_bound, + ) + scratch.ssm_state.normal_(0, 0.02) + copied = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + bench_reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + # 计时用均匀窗口。宽度为 1 时 verify_width 就是 1。 + cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) + cu = (cu * state.verify_width).clamp_max(tokens) + return (), dict( + state=scratch, + linear_layer_idx=0, + q=copied[0], + k=copied[1], + v=copied[2], + a=copied[3], + b=copied[4], + a_log=copied[5], + bias=copied[6], + reqs=bench_reqs, + cu_seqlens=cu, + ) + + +def _run_key_distance(left: str, right: str) -> int: + return sum(abs(x - y) for x, y in zip(ast.literal_eval(left), ast.literal_eval(right))) + + +@autotune( + kernel_name="replay_kda_decode:v1", + configs_gen_func=get_configs, + static_key_func=get_static_key, + run_key_func=get_run_key, + run_key_distance_func=_run_key_distance, + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) def replay( state, linear_layer_idx: int, @@ -25,31 +128,25 @@ def replay( reqs: torch.Tensor, cu_seqlens: torch.Tensor, run_config: dict[str, int] | None = None, -) -> tuple[torch.Tensor, dict[str, int]]: +) -> torch.Tensor: """执行一层 KDA replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 """ assert state.kda, "KDA replay step requires the KDA gate" - assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" - assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" - assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" - # 没传入时走 autotune 或默认 BV/warp/stage。传入的缺项也在这里补，不在 state 上改写。 - config = ( - _replay_config(run_config) - if run_config is not None - else resolve_run_config(state, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - ) + _check_replay_inputs(q, k, v, reqs, cu_seqlens) + if state._config_is_fixed and state.run_config is not None: + run_config = state.run_config + elif run_config is None: + run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) out = torch.empty_like(v) - bv = config["BV"] + bv = run_config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 precompute_projections = state.compute_mode == "precompute" and bv >= 16 - # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 padded_verify_width = triton.next_power_of_2(state.verify_width) - sync_mark = state.tile_sync_mark _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -67,7 +164,7 @@ def replay( state.betas[linear_layer_idx], reqs, state.cursors, - sync_mark, + state.tile_sync_mark, cu_seqlens, out, q.stride(0), @@ -91,7 +188,7 @@ def replay( BQK=max(16, 2 * padded_verify_width), BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), + num_warps=run_config["num_warps"], + num_stages=run_config["num_stages"], ) - return out, config + return out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py deleted file mode 100644 index ab29f3ddd4..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ /dev/null @@ -1,178 +0,0 @@ -"""使用共享 autotuner 调优单层 SSM update cycle。""" - -from __future__ import annotations - -import ast - -import torch -import triton - -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -_DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} - - -def resolve_run_config(state, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: - """返回当前 batch 的固定或已调优 SSM kernel 配置。""" - if state._config_is_fixed and state.run_config is not None: - return _replay_config(state.run_config) - default = _replay_config(state.run_config) - key = (tuple(static_key(state, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) - if key in state.batch_configs: - return state.batch_configs[key] - if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - return default - config = _replay_config(select_config(state, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) - # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 - state.batch_configs[key] = config - logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") - return config - - -def get_configs() -> list[dict[str, int]]: - """返回 ReplaySSM update cycle 的候选 Triton 配置。""" - return [ - {"BV": bv, "num_warps": warps, "num_stages": stages} - for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] - for stages in (1, 2, 3) - ] - - -def static_key(state, mode, q, cu_seqlens) -> dict[str, object]: - """构造区分状态布局、数据类型和投影模式的静态调优键。""" - _, _, hv, kd, vd = state.ssm_state.shape - # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 - key = { - "mode": mode, - "v": 10, - "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), - "dtype": str((q.dtype, state.ssm_state.dtype)), - "capacity": state.capacity, - "kda": int(state.kda), - } - if state.kda: - key["lower"] = state.gate_lower_bound - key["recon_dtype"] = str(state.keys.dtype) - key["projection"] = ("kda_mma" if state.kda else "gdn_mma") if state.compute_mode == "precompute" else "inline" - return key - - -def run_key(q, cu_seqlens) -> tuple[int, int]: - """返回调优键：token 数，以及请求数。q 是 [tokens, heads, dim]。没传前缀和时一行一个请求。""" - tokens = q.shape[0] - return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens - - -def rebuild_inputs( - state, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None -) -> tuple[tuple[object, ...], dict[str, object]]: - """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" - tokens, run_count = run_key(q, cu_seqlens) - # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 - batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) - _, _, hv, kd, vd = state.ssm_state.shape - from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig - - # 计时只需要一层、本 batch 的槽位。正式执行仍用调用方原来的 state。 - scratch = type(state)( - ConvSSMStateConfig( - logical_layer_ids=(0,), - global_k_heads=state.num_key_heads, - global_v_heads=hv, - key_dim=kd, - value_dim=vd, - conv_width=1, - conv_dtype=q.dtype, - ssm_dtype=state.ssm_state.dtype, - ), - batch + 1, - mtp_step=state.verify_width - 1, - device=q.device, - ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, - gate_lower_bound=state.gate_lower_bound, - ) - scratch.ssm_state.normal_(0, 0.02) - # benchmark 回调不能再次触发配置选择。 - scratch._config_is_fixed = True - inputs = [ - torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) - ] - reqs = torch.cat( - ( - torch.arange(batch, dtype=torch.int32, device=q.device), - torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), - ) - ) - # 计时用均匀窗口，不读调用方的前缀和。宽度为 1 时 verify_width 就是 1。 - cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) - cu = (cu * state.verify_width).clamp_max(tokens) - return (), dict( - state=scratch, - 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, - workload=(reqs,), - ) - - -def replay_layer(state, *args, **kwargs): - """按 state 的具体类型调用对应的 replay decode。""" - - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - - if isinstance(state, ReplayKDAState): - from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay - elif isinstance(state, ReplayGDNState): - from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay - else: - raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(state).__name__}") - return replay(state, *args, **kwargs) - - -@autotune( - kernel_name="ssm_update_cycle:v2", - 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( - state, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None -) -> dict[str, int] | None: - """搜索 update cycle 布局；计时只使用重建的临时输入。""" - if workload is None: - return run_config - state.run_config = run_config - (reqs,) = workload - state.cursors.zero_() - # 把前缀写满但不越出 capacity。replay kernel 自己读 cursor。调优窗口全部接受。 - tokens = run_key(q, cu_seqlens)[0] - rounds = max(1, state.capacity // min(state.verify_width, tokens)) - for _ in range(rounds): - replay_layer(state, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) - advance_cursor(state, reqs, cu_seqlens) - return run_config - - -def _replay_config(config: dict | None) -> dict[str, int]: - if not config: - return dict(_DEFAULT_REPLAY_CONFIG) - return { - "BV": int(config["BV"]), - "num_warps": int(config["num_warps"]), - "num_stages": int(config.get("num_stages", 3)), - } diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 39a3d1f388..fb6ed69772 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -14,7 +14,7 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer +from lightllm.common.basemodel.attention.linear.replay_kda_decode import get_configs, replay from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, @@ -74,7 +74,7 @@ def cycle(): num_accepted_tokens=counts if cu is not None else None, ) else: - replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, replay_cu) + replay(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, replay_cu) if ssm_backend is not None: accept_replay_window(ssm_backend, reqs, accepted, replay_cu) if isinstance(ssm_backend, ReplaySSMExecutor): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 84b6e8db99..ad0ec66a5e 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -51,7 +51,6 @@ def attach_replay_history( owner.compute_mode = compute_mode owner.run_config = run_config owner._config_is_fixed = run_config is not None - owner.batch_configs = {} layers, slots, hv, k, v = ssm.shape owner.num_key_heads = hv if num_key_heads is None else num_key_heads assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 @@ -92,6 +91,16 @@ def _trim_rejected_index( tl.store(cursors + r, cursor - (window - accepted), valid) +def replay_layer(state, *args, **kwargs): + """按测试 state 的类型调用对应的 replay decode。""" + + if state.kda: + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + else: + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + return replay(state, *args, **kwargs) + + def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens=None) -> None: """测试里一次结算 cursor：先按整窗全接受，再按槽位 index 剪掉未接受尾部。 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 index f228eb7133..ddb02125ca 100644 --- 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 @@ -8,6 +8,7 @@ from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, + replay_layer, ) @@ -26,7 +27,6 @@ def ssm_snapshot(backend, req): from types import SimpleNamespace from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer def make(n, width): @@ -198,7 +198,7 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput def cycle(i): ssm_backend = ssm_backends[i] - outputs = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + outputs = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] (accept_unfused if i == 0 else accept_fused)(managers[i], args) return outputs 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 index 894bcd5276..d77fc5b5d8 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,9 +1,11 @@ import pytest import torch +from lightllm.common.basemodel.attention.linear.replay_gdn_decode import get_configs from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, + replay_layer, ) @@ -23,7 +25,6 @@ def ssm_snapshot(backend, req): from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -59,7 +60,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state num_key_heads=1, kda=kda, compute_mode="precompute", - run_config={"BV": bv, "num_warps": 1}, + run_config={"BV": bv, "num_warps": 1, "num_stages": 3}, ) reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") @@ -82,7 +83,7 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state # an overflowed history correction by reconstructing from raw inputs. assert ssm_backend.cursors[0].item() == 2 * round_idx assert not state.any() - actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] + actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu) expected = [] for token in range(2): output, expected_state = reference( @@ -185,7 +186,7 @@ def test_replay_acceptance_flush_and_graph( accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): - out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] accept_replay_window(ssm_backend, reqs, accepted, cu) return out @@ -285,7 +286,7 @@ def test_precomputed_projections_match_independent_reference(capacity, width, dt (128, 128), capacity, dtype, - run_config={"BV": 32, "num_warps": 1}, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, compute_mode="precompute", ) @@ -314,7 +315,7 @@ def test_production_gdn_geometry_matches_independent_reference(dtype): 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} + config = {"BV": bv, "num_warps": warps, "num_stages": 3} ssm_backends = [ ReplaySSMExecutor( state.clone(), @@ -357,7 +358,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d for ssm_backend in ssm_backends: ssm_backend.fold_accepted_prefix(overflow_reqs) for ssm_backend in ssm_backends: - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)) accept_replay_window(ssm_backend, reqs, accepted, cu) if dtype == torch.float32: torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) @@ -394,7 +395,7 @@ def test_bf16_commit_rounds_before_the_next_token(): cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") for ssm_backend in ssm_backends: assert ssm_backend.cursors[0].item() == 0 - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0]) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)) torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) @@ -432,7 +433,7 @@ def check(capacity, width, steps): v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 b = torch.randn_like(a) - outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] for backend in backends] + outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) for backend in backends] # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) for backend in backends: @@ -476,8 +477,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): assert int(backend.cursors[0]) == capacity outs, twin_outs = [], [] for layer in range(2): - outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0]) - twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu)[0]) + outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu)) + twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu)) if step == capacity: assert int(backend.cursors[0]) == capacity for actual, expected in zip(outs, twin_outs): @@ -495,7 +496,7 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): state = torch.randn(1, 3, 2, 32, 32, device="cuda") backend = ReplaySSMExecutor( - state, 8, 4, torch.float32, num_key_heads=2, kda=kda, run_config={"BV": 32, "num_warps": 1} + state, 8, 4, torch.float32, num_key_heads=2, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} ) backend.cursors[0] = 3 reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) @@ -527,7 +528,7 @@ def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_ num_key_heads=2, kda=kda, compute_mode=compute_mode, - run_config={"BV": 32, "num_warps": 1}, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ) expected = state.float().clone() pending = [0, 0] @@ -543,7 +544,7 @@ def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_ bias = torch.zeros(4 * 128 if kda else 4, device="cuda") def forward(): - return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] def step(): positions = ssm_backend.cursors.index_select(0, reqs.long()) @@ -750,7 +751,9 @@ def test_kda_tuning_candidates_match_reference(config, compute_mode): def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): torch.manual_seed(72) initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 - backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) + backend = ReplaySSMExecutor( + initial.clone(), 8, 4, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} + ) reqs = torch.cat( ( torch.arange(live, device="cuda", dtype=torch.int32), @@ -784,7 +787,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa bias, reqs, cu, - )[0].reshape(live, query_width, 1, 32) + ).reshape(live, query_width, 1, 32) for token in range(query_width): output, expected = reference( q[:, token], @@ -815,7 +818,14 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): dtype = torch.bfloat16 state = torch.randn(1, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 backend = ReplaySSMExecutor( - state, 4, 2, dtype, num_key_heads=kh, kda=kda, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + state, + 4, + 2, + dtype, + num_key_heads=kh, + kda=kda, + compute_mode=compute_mode, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ) reqs = torch.tensor([0], device="cuda", dtype=torch.int32) indices = torch.zeros(2, device="cuda", dtype=torch.int32) 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 index e4693181e0..ef4a1c618c 100644 --- 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 @@ -4,8 +4,8 @@ from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, + replay_layer, ) -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) @@ -37,7 +37,9 @@ def record(*args, run=kernel.run, name=name, **kwargs): torch.manual_seed(137) initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 eager, captured = [ - ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1}) + ReplaySSMExecutor( + initial.clone(), 8, 1, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} + ) for _ in range(2) ] reqs = torch.arange(512, device="cuda", dtype=torch.int32) @@ -50,7 +52,7 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) - out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids, cu)[0] + out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids, cu) accept_replay_window(backend, ids, cu_seqlens=cu) return out @@ -153,7 +155,7 @@ def metadata(ssm_backend): def forward_accept(index, cu, reqs): ssm_backend = ssm_backends[index] - out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] if index == 0: accept_replay_window(ssm_backend, reqs, accepted, cu) else: @@ -263,7 +265,7 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) def forward(ssm_backend): - return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] + return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] def warm(ssm_backend): out = forward(ssm_backend) @@ -352,7 +354,7 @@ def step(ssm_backend, accepted): if overflow: ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) pos = ssm_backend.cursors.index_select(0, reqs.long()) - out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(48)] accept_replay_window(ssm_backend, reqs, accepted, cu) return pos, out @@ -400,9 +402,7 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): def step(): - out = [ - replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2) - ] + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets) for layer in range(2)] accept_replay_window(ssm_backend, ids, accepted, offsets) return out 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 index b6c0aca7b0..a94411fbd5 100644 --- 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 @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( attach_replay_history as ReplaySSMExecutor, ) @@ -17,7 +17,8 @@ def test_rebuild_uses_one_layer_scratch(width): gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") a_log = torch.zeros(2, device="cuda") cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + _, kwargs = tuning.rebuild_inputs(ssm_backend, 0, q, q, q, gates, gates, a_log, a_log, reqs, cu) scratch = kwargs["state"] assert scratch is not ssm_backend assert scratch.ssm_state.shape[0] == 1 @@ -42,15 +43,13 @@ def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, w a_log = torch.zeros(2, device="cuda") bias = torch.zeros(2, device="cuda") cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - inputs = dict( - state=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu - ) + reqs = torch.arange(3, device="cuda", dtype=torch.int32) before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} def oom(*args, **kwargs): raise torch.cuda.OutOfMemoryError("scratch") - tuner = tuning.select_config + tuner = tuning.replay monkeypatch.setattr(tuner, "rebuild_input_func", oom) monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) monkeypatch.setattr(tuner, "cached_configs", {}) @@ -61,12 +60,11 @@ def oom(*args, **kwargs): ) with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - tuning.resolve_run_config(**inputs) + tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu) assert not list(tmp_path.iterdir()) config = {"BV": 32, "num_warps": 1, "num_stages": 3} for name, expected in before.items(): torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) - reqs = torch.arange(3, device="cuda", dtype=torch.int32) - tuning.replay_layer(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) torch.cuda.synchronize() diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index cf89b94480..14f6c7cb7d 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -5,7 +5,6 @@ from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -70,17 +69,6 @@ def test_kda_backend_receives_gate_bound_at_construction(linear_config): assert owner.gate_lower_bound == -2.5 -def test_replay_backend_publishes_detached_run_config(): - ssm_backend = object.__new__(ReplayGDNState) - tuning = {"BV": 32, "num_warps": 2} - ssm_backend.run_config = tuning - ssm_backend._config_is_fixed = True - ssm_backend.compute_mode = "inline" - config = resolve_run_config(ssm_backend, "replay", *([None] * 7), None) - tuning["BV"] = 128 - assert config == {"BV": 32, "num_warps": 2, "num_stages": 3} - - @pytest.mark.parametrize("cls", [GDNState, KDAState]) @pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): From bc474ff84aad22de7e6cf7815c3cab8997e4fcd6 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 04:36:38 +0000 Subject: [PATCH 45/54] refactor: pass replay tile configs as arguments A pinned BV and warp count is an input to the launch. Serving state no longer carries a flag that overrides the autotuner. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 2 -- .../bounded_att/kda/replay_kda_state.py | 2 -- .../attention/linear/replay_gdn_decode.py | 4 +--- .../attention/linear/replay_kda_decode.py | 4 +--- test/benchmark/kernels/benchmark_kda_ssm.py | 17 ++++++++++++++--- .../triton_kernel/linear_att/replay_history.py | 3 ++- 6 files changed, 18 insertions(+), 14 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index f972d4b884..ab2bfaf490 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -94,8 +94,6 @@ def __init__( self.capacity = capacity self.verify_width = verify_width self.compute_mode = compute_mode - self.run_config = None - self._config_is_fixed = False layers, slots, hv, k, v = self.ssm_state.shape self.num_key_heads = config.num_k_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 133d6063e8..cbb068b164 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -94,8 +94,6 @@ def __init__( self.capacity = capacity self.verify_width = verify_width self.compute_mode = compute_mode - self.run_config = None - self._config_is_fixed = False layers, slots, hv, k, v = self.ssm_state.shape self.num_key_heads = config.num_k_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index fc345fe34e..46a5b09795 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -135,9 +135,7 @@ def replay( assert not state.kda, "GDN replay step does not use the KDA gate" _check_replay_inputs(q, k, v, reqs, cu_seqlens) - if state._config_is_fixed and state.run_config is not None: - run_config = state.run_config - elif run_config is None: + if run_config is None: run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 50dc383f6f..3a59a16120 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -136,9 +136,7 @@ def replay( assert state.kda, "KDA replay step requires the KDA gate" _check_replay_inputs(q, k, v, reqs, cu_seqlens) - if state._config_is_fixed and state.run_config is not None: - run_config = state.run_config - elif run_config is None: + if run_config is None: run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} hv, kd, vd = state.ssm_state.shape[-3:] assert state.num_key_heads in (q.shape[-2], hv) diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index fb6ed69772..de48a1e381 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -43,8 +43,6 @@ def workload(args, batch, width, dtype, mode, acceptance): if mode != "native": capacity = max(8, triton.next_power_of_2(width)) ssm_backend = ReplaySSMExecutor(state, capacity, width, q.dtype, kda=True, compute_mode=mode) - 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 @@ -74,7 +72,20 @@ def cycle(): num_accepted_tokens=counts if cu is not None else None, ) else: - replay(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, replay_cu) + replay( + ssm_backend, + layer, + q, + k, + v, + gate, + beta, + alog, + bias, + reqs, + replay_cu, + run_config=ssm_backend.run_config, + ) if ssm_backend is not None: accept_replay_window(ssm_backend, reqs, accepted, replay_cu) if isinstance(ssm_backend, ReplaySSMExecutor): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index ad0ec66a5e..3d2eb7027d 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -50,7 +50,6 @@ def attach_replay_history( owner.verify_width = verify_width owner.compute_mode = compute_mode owner.run_config = run_config - owner._config_is_fixed = run_config is not None layers, slots, hv, k, v = ssm.shape owner.num_key_heads = hv if num_key_heads is None else num_key_heads assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 @@ -98,6 +97,8 @@ def replay_layer(state, *args, **kwargs): from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay else: from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + if "run_config" not in kwargs and state.run_config is not None: + kwargs["run_config"] = state.run_config return replay(state, *args, **kwargs) From ba78baa925a10ddd0638cf9dd397c6ddd47a543c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 04:58:20 +0000 Subject: [PATCH 46/54] refactor: give GDN and KDA separate linear backend entries. Each algorithm keeps its own native and replay class tables, and the chunk-index host sync is marked for removal. Co-authored-by: Cursor --- .../attention/create_linear_utils.py | 93 +++++++++++-------- .../basemodel/attention/linear/replay_gdn.py | 13 --- .../triton_kernel/linear_att/fla/ops/index.py | 1 + lightllm/models/glm5_next/model.py | 12 +-- lightllm/models/qwen3next/model.py | 8 +- .../basemodel/attention/linear/test_gdn.py | 51 +++++++--- .../basemodel/attention/test_lazy_import.py | 30 +++--- 7 files changed, 116 insertions(+), 92 deletions(-) diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index fce8e8d78f..08528f98f6 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -1,66 +1,77 @@ -from .import_utils import load_backend_class +from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend +from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend +from lightllm.common.basemodel.attention.linear.replay_gdn import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, +) +from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend from lightllm.utils.backend_validator import validate from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) -linear_prefill_att_backend_classes = { - "flashqla": "linear.flashqla:FlashQlaLinearAttBackend", - "triton": "linear.triton:TritonLinearAttBackend", +gdn_prefill_backends = { + "flashqla": FlashQlaLinearAttBackend, + "triton": TritonLinearAttBackend, } -linear_decode_att_backend_classes = { - "triton": "linear.triton:TritonLinearAttBackend", +gdn_replay_prefill_backends = { + "flashqla": ReplayGDNFlashQlaAttBackend, + "triton": ReplayGDNTritonAttBackend, } +gdn_decode_backends = { + "triton": TritonLinearAttBackend, +} -def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): - args = get_env_start_args() - backend_str = _get_backend_str(args.llm_prefill_att_backend, index) - if backend_str != "auto": - return _linear_backend_class(linear_prefill_att_backend_classes[backend_str], backend_str) +gdn_replay_decode_backends = { + "triton": ReplayGDNTritonAttBackend, +} - for backend_name in priority_list: - backend_class = linear_prefill_att_backend_classes[backend_name] - if backend_name == "triton": - logger.info("Linear prefill attention backend: triton.") - return _linear_backend_class(backend_class, backend_name) - if validate(backend_name): - logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return _linear_backend_class(backend_class, backend_name) +kda_backends = { + "triton": KDALinearAttBackend, +} - logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") +kda_replay_backends = { + "triton": ReplayKDALinearAttBackend, +} -def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): +def get_gdn_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): args = get_env_start_args() - backend_str = _get_backend_str(args.llm_decode_att_backend, index) - if backend_str != "auto": - return _linear_backend_class(linear_decode_att_backend_classes[backend_str], backend_str) + classes = gdn_replay_prefill_backends if args.ssm_state_mode == "replay" else gdn_prefill_backends + return _select_backend(classes, args.llm_prefill_att_backend, index, priority_list, "prefill") - for backend_name in priority_list: - backend_class = linear_decode_att_backend_classes[backend_name] - if backend_name == "triton": - logger.info("Linear decode attention backend: triton.") - return _linear_backend_class(backend_class, backend_name) - if validate(backend_name): - logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return _linear_backend_class(backend_class, backend_name) - logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") +def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)): + args = get_env_start_args() + classes = gdn_replay_decode_backends if args.ssm_state_mode == "replay" else gdn_decode_backends + return _select_backend(classes, args.llm_decode_att_backend, index, priority_list, "decode") -def _linear_backend_class(native_path: str, backend_name: str): - """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" +def get_kda_linear_att_backend_class(): + args = get_env_start_args() + classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends + return classes["triton"] - if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": - return load_backend_class(native_path) - from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class - return replay_gdn_backend_class(backend_name) +def _select_backend(classes, backend_args, index, priority_list, stage): + backend_str = _get_backend_str(backend_args, index) + if backend_str != "auto": + return classes[backend_str] + + for backend_name in priority_list: + if backend_name == "triton" or validate(backend_name): + if backend_name == "triton": + logger.info(f"Linear {stage} attention backend: triton.") + else: + logger.info(f"Linear {stage} attention backend: {backend_name} (validated).") + return classes[backend_name] + + logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.") + return classes["triton"] def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 5b96d21596..8e4a081f73 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -44,19 +44,6 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo -def replay_gdn_backend_class(name: str): - """按已选中的 linear kernel 名字返回对应的 replay backend。""" - - classes = { - "triton": ReplayGDNTritonAttBackend, - "flashqla": ReplayGDNFlashQlaAttBackend, - } - try: - return classes[name] - except KeyError: - raise ValueError(f"replay GDN has no backend named {name!r}") from None - - class ReplayGDNTritonAttBackend(BaseAttBackend): """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py index 8b1d59fc63..878e2296f4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py @@ -21,6 +21,7 @@ def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: @tensor_cache def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: + # TODO: tolist() 会把 chunk 数同步回 CPU。以后改成在设备上构造索引，去掉这次同步。 indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 60dfd016c8..3f20ad97ab 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -6,7 +6,6 @@ import triton from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager @@ -206,14 +205,11 @@ def _init_att_backend(self): self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): - if getattr(self.args, "ssm_state_mode", "native") == "replay": - from lightllm.common.basemodel.attention.linear.replay_kda import ( - ReplayKDALinearAttBackend, - ) + from lightllm.common.basemodel.attention.create_linear_utils import ( + get_kda_linear_att_backend_class, + ) - backend_cls = ReplayKDALinearAttBackend - else: - backend_cls = KDALinearAttBackend + backend_cls = get_kda_linear_att_backend_class() self.prefill_att_backend1 = backend_cls(model=self) self.decode_att_backend1 = self.prefill_att_backend1 diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 9da87c6acb..57de61241b 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -144,12 +144,12 @@ def _init_att_backend1(self): self.decode_att_backend1 = None return from lightllm.common.basemodel.attention.create_linear_utils import ( - get_qwen35_linear_prefill_att_backend_class, - get_qwen35_linear_decode_att_backend_class, + get_gdn_linear_prefill_att_backend_class, + get_gdn_linear_decode_att_backend_class, ) - linear_prefill_att_backend_class = get_qwen35_linear_prefill_att_backend_class(index=1) - linear_decode_att_backend_class = get_qwen35_linear_decode_att_backend_class(index=1) + linear_prefill_att_backend_class = get_gdn_linear_prefill_att_backend_class(index=1) + linear_decode_att_backend_class = get_gdn_linear_decode_att_backend_class(index=1) self.prefill_att_backend1 = linear_prefill_att_backend_class(model=self) self.decode_att_backend1 = linear_decode_att_backend_class(model=self) return diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index b61b9c02b3..344bfea57f 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -154,6 +154,7 @@ def auto_linear_backend_args(monkeypatch): lambda: SimpleNamespace( llm_prefill_att_backend=["fa3"], llm_decode_att_backend=["flashinfer"], + ssm_state_mode="native", ), ) @@ -169,7 +170,7 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear lambda name: validate_calls.append(name) or True, ) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) assert backend_class is FlashQlaLinearAttBackend backend = object.__new__(backend_class) @@ -194,8 +195,8 @@ def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): ReplayGDNTritonAttBackend, ) - prefill_backend = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) - decode_backend = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) assert prefill_backend is ReplayGDNFlashQlaAttBackend assert decode_backend is ReplayGDNTritonAttBackend @@ -205,6 +206,25 @@ def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): assert not issubclass(decode_backend, LinearAttBackend) +def test_kda_entry_selects_native_or_replay(monkeypatch): + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="native"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="replay"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend + + def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, @@ -274,6 +294,7 @@ def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): lambda: SimpleNamespace( llm_prefill_att_backend=["fa3", "triton"], llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="native", ), ) monkeypatch.setattr( @@ -282,8 +303,8 @@ def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), ) - prefill_backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) - decode_backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) assert prefill_backend_class is TritonLinearAttBackend assert decode_backend_class is TritonLinearAttBackend @@ -308,8 +329,8 @@ def test_cli_accepts_linear_backend_at_index_one(): @pytest.mark.parametrize( "get_backend_class", [ - linear_create_utils.get_qwen35_linear_prefill_att_backend_class, - linear_create_utils.get_qwen35_linear_decode_att_backend_class, + linear_create_utils.get_gdn_linear_prefill_att_backend_class, + linear_create_utils.get_gdn_linear_decode_att_backend_class, ], ) def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): @@ -318,7 +339,7 @@ def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_b def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): - backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) assert backend_class is TritonLinearAttBackend @@ -327,8 +348,8 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ validate_calls = [] monkeypatch.setattr( linear_create_utils, - "linear_decode_att_backend_classes", - {"optimized": "linear.flashqla:FlashQlaLinearAttBackend", "triton": "linear.triton:TritonLinearAttBackend"}, + "gdn_decode_backends", + {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, ) monkeypatch.setattr( linear_create_utils, @@ -336,7 +357,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ lambda name: validate_calls.append(name) or True, ) - backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class( + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( index=1, priority_list=("optimized", "triton") ) @@ -347,7 +368,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) assert backend_class is TritonLinearAttBackend @@ -356,8 +377,8 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ validate_calls = [] monkeypatch.setattr( linear_create_utils, - "linear_prefill_att_backend_classes", - {"flashqla2": "linear.triton:TritonLinearAttBackend", "flashqla3": "linear.flashqla:FlashQlaLinearAttBackend"}, + "gdn_prefill_backends", + {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, ) monkeypatch.setattr( linear_create_utils, @@ -365,7 +386,7 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ lambda name, *args: validate_calls.append(name) or name == "flashqla3", ) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class( + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( index=1, priority_list=("flashqla2", "flashqla3") ) diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index 1efed8423e..e7423d30f9 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -67,6 +67,7 @@ def setUp(self): llm_decode_att_backend=["triton"], enable_ep_moe=False, mtp_step=0, + ssm_state_mode="native", ) self.validate = Mock(return_value=True) dependencies = { @@ -89,7 +90,6 @@ def setUp(self): self.attention = importlib.import_module(ATTENTION) def test_importing_selection_apis_loads_no_implementations(self): - importlib.import_module(f"{ATTENTION}.create_linear_utils") importlib.import_module(f"{ATTENTION}.linear") importlib.import_module(f"{ATTENTION}.nsa") self.assertEqual(self.imports.imported, []) @@ -184,24 +184,32 @@ def test_neo_auto_selection_uses_supported_fa3(self): self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) self.validate.assert_called_once_with("fa3") - def test_linear_selection_imports_only_validated_backend(self): - backend = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + def _allow_linear_backend_classes(self): + flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") + self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") + replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] + replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) + self.imports.allow("linear.kda", "KDALinearAttBackend") + self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") + return flashqla, triton + + def test_linear_prefill_auto_selects_validated_flashqla(self): + flashqla, _triton = self._allow_linear_backend_classes() selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.flashqla"]) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) self.validate.assert_called_once_with("flashqla") - def test_linear_selection_falls_back_without_importing_flashqla(self): - backend = self.imports.allow("linear.triton", "TritonLinearAttBackend") + def test_linear_selection_falls_back_to_triton(self): + _flashqla, triton = self._allow_linear_backend_classes() selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") self.validate.return_value = False - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertIs(selection.get_qwen35_linear_decode_att_backend_class(), backend) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) + self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) self.args.llm_prefill_att_backend.append("triton") - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.triton"]) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) def test_vit_selection_loads_only_requested_implementation(self): package = "lightllm.common.basemodel.attention_vit" From 85f0f2219154660476e297bd9174acea1fb3f14f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 05:25:52 +0000 Subject: [PATCH 47/54] refactor: pass KDA replay tensors instead of the state object. The decode kernel receives the layer SSM, history, and cursor directly, so the operator no longer indexes the serving state. Co-authored-by: Cursor --- .../basemodel/attention/linear/replay_kda.py | 33 ++- .../attention/linear/replay_kda_decode.py | 209 +++++++++++------- test/benchmark/kernels/benchmark_kda_ssm.py | 32 ++- .../linear_att/replay_history.py | 34 ++- .../common/test_linear_state_contract.py | 64 ++++-- 5 files changed, 245 insertions(+), 127 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index e5e8d4df85..af931aa477 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -182,18 +182,29 @@ def decode_att( v = v.view(-1, backend.tp_num_heads, backend.head_dim) raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) + layer = state.storage_layer(layer_num) output = replay( - state, - state.storage_layer(layer_num), - q, - k, - v, - raw_gate, - raw_beta, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b_cu_seqlens, + state_ssm=state.ssm_state[layer], + state_keys=state.keys[layer], + state_deltas=state.deltas[layer], + state_gates=state.gates[layer], + state_raw_keys=state.raw_keys[layer], + state_raw_values=state.raw_values[layer], + state_betas=state.betas[layer], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=raw_gate, + b=raw_beta, + a_log=layer_weight.linear_A_log.weight, + bias=layer_weight.linear_dt_bias.weight, + reqs=self.b_conv_buffer_idx, + cu_seqlens=self.b_cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, ) # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 3a59a16120..eedbec91b7 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -1,6 +1,6 @@ """KDA replay decode。 -History 张量和 cursor 在 ReplayKDAState。这里调优并执行一层重放。 +调用方传入这一层的 SSM、history 和 cursor。这里调优并执行一层重放。 """ from __future__ import annotations @@ -10,7 +10,6 @@ import torch import triton -from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune @@ -24,17 +23,24 @@ def get_configs() -> list[dict[str, int]]: ] -def get_static_key(state, q) -> dict[str, object]: +def get_static_key( + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + q: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, +) -> dict[str, object]: """区分头维度、状态精度、history 容量、gate 下界和投影模式。""" - _, _, hv, kd, vd = state.ssm_state.shape + _, hv, kd, vd = state_ssm.shape return { "v": 1, - "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), - "dtype": str((q.dtype, state.ssm_state.dtype)), - "capacity": state.capacity, - "lower": state.gate_lower_bound, - "recon_dtype": str(state.keys.dtype), - "projection": "kda_mma" if state.compute_mode == "precompute" else "inline", + "shape": str((q.shape[-2], hv, state_keys.shape[-3], kd, vd, state_verify_width)), + "dtype": str((q.dtype, state_ssm.dtype)), + "capacity": state_keys.shape[-2], + "lower": state_gate_lower_bound, + "recon_dtype": str(state_keys.dtype), + "projection": "kda_mma" if state_compute_mode == "precompute" else "inline", } @@ -49,46 +55,62 @@ def _check_replay_inputs(q, k, v, reqs, cu_seqlens) -> None: assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" -def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens, run_config=None): +def rebuild_inputs( + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + state_deltas: torch.Tensor, + state_gates: torch.Tensor, + state_raw_keys: torch.Tensor, + state_raw_values: torch.Tensor, + state_betas: torch.Tensor, + state_cursors: torch.Tensor, + state_tile_sync_mark: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, + run_config: dict[str, int] | None = None, +) -> tuple[tuple, dict[str, object]]: """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" - del linear_layer_idx, reqs, run_config + del state_cursors, reqs, run_config tokens, run_count = get_run_key(q, cu_seqlens) # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 - batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) - _, _, hv, kd, vd = state.ssm_state.shape - scratch = type(state)( - ConvSSMStateConfig( - logical_layer_ids=(0,), - global_k_heads=state.num_key_heads, - global_v_heads=hv, - key_dim=kd, - value_dim=vd, - conv_width=1, - conv_dtype=q.dtype, - ssm_dtype=state.ssm_state.dtype, - ), - batch + 1, - mtp_step=state.verify_width - 1, - device=q.device, - ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, - gate_lower_bound=state.gate_lower_bound, - ) - scratch.ssm_state.normal_(0, 0.02) + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state_verify_width)) + slots = batch + 1 + _, hv, kd, vd = state_ssm.shape + key_heads, capacity = state_keys.shape[-3], state_keys.shape[-2] + device = q.device + rows_per_slot = state_tile_sync_mark.shape[0] // state_ssm.shape[0] copied = [ torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) ] bench_reqs = torch.cat( ( - torch.arange(batch, dtype=torch.int32, device=q.device), - torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), + torch.arange(batch, dtype=torch.int32, device=device), + torch.full((run_count - batch,), batch, dtype=torch.int32, device=device), ) ) # 计时用均匀窗口。宽度为 1 时 verify_width 就是 1。 - cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) - cu = (cu * state.verify_width).clamp_max(tokens) + cu = torch.arange(run_count + 1, dtype=torch.int32, device=device) + cu = (cu * state_verify_width).clamp_max(tokens) return (), dict( - state=scratch, - linear_layer_idx=0, + state_ssm=torch.empty(slots, hv, kd, vd, device=device, dtype=state_ssm.dtype).normal_(0, 0.02), + state_keys=torch.empty(slots, key_heads, capacity, kd, device=device, dtype=state_keys.dtype), + state_deltas=torch.empty(slots, hv, capacity, vd, device=device, dtype=state_deltas.dtype), + state_gates=torch.empty(slots, hv, capacity, kd, device=device, dtype=state_gates.dtype), + state_raw_keys=torch.empty(slots, key_heads, capacity, kd, device=device, dtype=state_raw_keys.dtype), + state_raw_values=torch.empty(slots, hv, capacity, vd, device=device, dtype=state_raw_values.dtype), + state_betas=torch.empty(slots, hv, capacity, device=device, dtype=state_betas.dtype), + state_cursors=torch.zeros(slots, dtype=torch.int32, device=device), + state_tile_sync_mark=torch.zeros((slots * rows_per_slot, hv), dtype=torch.int32, device=device), q=copied[0], k=copied[1], v=copied[2], @@ -98,6 +120,9 @@ def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu bias=copied[6], reqs=bench_reqs, cu_seqlens=cu, + state_verify_width=state_verify_width, + state_compute_mode=state_compute_mode, + state_gate_lower_bound=state_gate_lower_bound, ) @@ -116,8 +141,15 @@ def _run_key_distance(left: str, right: str) -> int: warmup_all_exist_config=False, ) def replay( - state, - linear_layer_idx: int, + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + state_deltas: torch.Tensor, + state_gates: torch.Tensor, + state_raw_keys: torch.Tensor, + state_raw_values: torch.Tensor, + state_betas: torch.Tensor, + state_cursors: torch.Tensor, + state_tile_sync_mark: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -127,64 +159,71 @@ def replay( bias: torch.Tensor, reqs: torch.Tensor, cu_seqlens: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, run_config: dict[str, int] | None = None, ) -> torch.Tensor: - """执行一层 KDA replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 + """执行一层 KDA replay decode。 - q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 + state_* 张量已经是这一层的视图。ssm 是 [slots, value_heads, key_dim, value_dim]， + keys 是 [slots, key_heads, capacity, key_dim]。最后一个 slot 是 HOLD。 + q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。 + 宽度为 1 时它是 0, 1, 2, ...。 """ - assert state.kda, "KDA replay step requires the KDA gate" _check_replay_inputs(q, k, v, reqs, cu_seqlens) if run_config is None: run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} - hv, kd, vd = state.ssm_state.shape[-3:] - assert state.num_key_heads in (q.shape[-2], hv) + _, hv, kd, vd = state_ssm.shape + key_heads = state_keys.shape[-3] + capacity = state_keys.shape[-2] + assert key_heads in (q.shape[-2], hv) out = torch.empty_like(v) bv = run_config["BV"] # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 - precompute_projections = state.compute_mode == "precompute" and bv >= 16 - padded_verify_width = triton.next_power_of_2(state.verify_width) + precompute_projections = state_compute_mode == "precompute" and bv >= 16 + padded_verify_width = triton.next_power_of_2(state_verify_width) _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( - q, - k, - v, - a, - b, - a_log, - bias, - state.ssm_state[linear_layer_idx], - state.keys[linear_layer_idx], - state.deltas[linear_layer_idx], - state.gates[linear_layer_idx], - state.raw_keys[linear_layer_idx], - state.raw_values[linear_layer_idx], - state.betas[linear_layer_idx], - reqs, - state.cursors, - state.tile_sync_mark, - cu_seqlens, - out, - q.stride(0), - k.stride(0), - v.stride(0), - a.stride(0), - b.stride(0), - q.shape[-2], - hv, - state.num_key_heads, - kd, - vd, - state.capacity, - state.hold_req_id, - triton.next_power_of_2(kd), - bv, - precompute_projections, - padded_verify_width, - True, - state.gate_lower_bound, + Q=q, + StepK=k, + StepV=v, + A=a, + B=b, + Alog=a_log, + Bias=bias, + State=state_ssm, + Keys=state_keys, + Deltas=state_deltas, + Gates=state_gates, + RawKeys=state_raw_keys, + RawValues=state_raw_values, + Betas=state_betas, + Reqs=reqs, + Cursors=state_cursors, + SyncMark=state_tile_sync_mark, + Cu=cu_seqlens, + Out=out, + SQ=q.stride(0), + SK=k.stride(0), + SV=v.stride(0), + SA=a.stride(0), + SB=b.stride(0), + H=q.shape[-2], + HV=hv, + KH=key_heads, + K=kd, + V=vd, + L=capacity, + HOLD=state_ssm.shape[0] - 1, + BK=triton.next_power_of_2(kd), + BV=bv, + PRECOMPUTE_PROJECTIONS=precompute_projections, + BW=padded_verify_width, + KDA=True, + LOWER=state_gate_lower_bound, BQK=max(16, 2 * padded_verify_width), - BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), + BR=triton.next_power_of_2(max(16, capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), num_warps=run_config["num_warps"], num_stages=run_config["num_stages"], diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index de48a1e381..bd12f94321 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -73,17 +73,27 @@ def cycle(): ) else: replay( - ssm_backend, - layer, - q, - k, - v, - gate, - beta, - alog, - bias, - reqs, - replay_cu, + state_ssm=ssm_backend.ssm_state[layer], + state_keys=ssm_backend.keys[layer], + state_deltas=ssm_backend.deltas[layer], + state_gates=ssm_backend.gates[layer], + state_raw_keys=ssm_backend.raw_keys[layer], + state_raw_values=ssm_backend.raw_values[layer], + state_betas=ssm_backend.betas[layer], + state_cursors=ssm_backend.cursors, + state_tile_sync_mark=ssm_backend.tile_sync_mark, + q=q, + k=k, + v=v, + a=gate, + b=beta, + a_log=alog, + bias=bias, + reqs=reqs, + cu_seqlens=replay_cu, + state_verify_width=ssm_backend.verify_width, + state_compute_mode=ssm_backend.compute_mode, + state_gate_lower_bound=ssm_backend.gate_lower_bound, run_config=ssm_backend.run_config, ) if ssm_backend is not None: diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 3d2eb7027d..e6fdd1fda3 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -93,12 +93,38 @@ def _trim_rejected_index( def replay_layer(state, *args, **kwargs): """按测试 state 的类型调用对应的 replay decode。""" - if state.kda: - from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay - else: - from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay if "run_config" not in kwargs and state.run_config is not None: kwargs["run_config"] = state.run_config + if state.kda: + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + + layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens = args + return replay( + state_ssm=state.ssm_state[layer_idx], + state_keys=state.keys[layer_idx], + state_deltas=state.deltas[layer_idx], + state_gates=state.gates[layer_idx], + state_raw_keys=state.raw_keys[layer_idx], + state_raw_values=state.raw_values[layer_idx], + state_betas=state.betas[layer_idx], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=a_log, + bias=bias, + reqs=reqs, + cu_seqlens=cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + run_config=kwargs.get("run_config"), + ) + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + return replay(state, *args, **kwargs) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 14f6c7cb7d..33c4b6be23 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -101,6 +101,7 @@ def __init__(self): def __getitem__(self, grid): def launch(*args, **kwargs): self.args = args + self.kwargs = kwargs return launch @@ -131,22 +132,53 @@ def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, modul q = torch.zeros(1, heads, linear_config.key_dim) v = torch.zeros(1, owner.ssm_state.shape[-3], linear_config.value_dim) gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] - decode.replay( - owner, - 0, - q, - q, - v, - torch.zeros(1, gate), - torch.zeros(1, v.shape[-3]), - torch.zeros(v.shape[-3]), - torch.zeros(gate), - torch.zeros(1, dtype=torch.int32), - torch.arange(2, dtype=torch.int32), - run_config={"BV": bv, "num_warps": 1, "num_stages": 1}, - ) - assert captured.args[33] is precompute - assert captured.args[35] is (cls is ReplayKDAState) + run_config = {"BV": bv, "num_warps": 1, "num_stages": 1} + if cls is ReplayKDAState: + decode.replay( + state_ssm=owner.ssm_state[0], + state_keys=owner.keys[0], + state_deltas=owner.deltas[0], + state_gates=owner.gates[0], + state_raw_keys=owner.raw_keys[0], + state_raw_values=owner.raw_values[0], + state_betas=owner.betas[0], + state_cursors=owner.cursors, + state_tile_sync_mark=owner.tile_sync_mark, + q=q, + k=q, + v=v, + a=torch.zeros(1, gate), + b=torch.zeros(1, v.shape[-3]), + a_log=torch.zeros(v.shape[-3]), + bias=torch.zeros(gate), + reqs=torch.zeros(1, dtype=torch.int32), + cu_seqlens=torch.arange(2, dtype=torch.int32), + state_verify_width=owner.verify_width, + state_compute_mode=owner.compute_mode, + state_gate_lower_bound=owner.gate_lower_bound, + run_config=run_config, + ) + else: + decode.replay( + owner, + 0, + q, + q, + v, + torch.zeros(1, gate), + torch.zeros(1, v.shape[-3]), + torch.zeros(v.shape[-3]), + torch.zeros(gate), + torch.zeros(1, dtype=torch.int32), + torch.arange(2, dtype=torch.int32), + run_config=run_config, + ) + if cls is ReplayKDAState: + assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute + assert captured.kwargs["KDA"] is True + else: + assert captured.args[33] is precompute + assert captured.args[35] is False @pytest.mark.parametrize("cls", [GDNState, KDAState]) From c79a062c28d3ed4ec27d4f3286af9e08f2814862 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 05:29:02 +0000 Subject: [PATCH 48/54] chore: record Cursor rules for new code. New calls use named arguments and type annotations, and production paths stay free of test-only APIs. Co-authored-by: Cursor --- .cursor/rules/named-args-and-types.mdc | 12 ++++++++++++ .cursor/rules/no-test-only-production-api.mdc | 10 ++++++++++ 2 files changed, 22 insertions(+) create mode 100644 .cursor/rules/named-args-and-types.mdc create mode 100644 .cursor/rules/no-test-only-production-api.mdc diff --git a/.cursor/rules/named-args-and-types.mdc b/.cursor/rules/named-args-and-types.mdc new file mode 100644 index 0000000000..a463ccac7d --- /dev/null +++ b/.cursor/rules/named-args-and-types.mdc @@ -0,0 +1,12 @@ +--- +description: 新代码的函数调用使用参数名，参数带类型标注 +alwaysApply: true +--- + +# 新代码的调用和类型 + +新写的函数调用使用 `param_name=value`。不要用 `*args`、`*activations`、`**kwargs` 把一组参数展开传进去。 + +新写的函数参数带类型标注。 + +修改旧代码时，原有调用和原有签名保持不动，避免把风格调整混进 review。这次改动里新加的函数、新加的调用按上面两条写。 diff --git a/.cursor/rules/no-test-only-production-api.mdc b/.cursor/rules/no-test-only-production-api.mdc new file mode 100644 index 0000000000..a9e877828e --- /dev/null +++ b/.cursor/rules/no-test-only-production-api.mdc @@ -0,0 +1,10 @@ +--- +description: 测试专用的实现和接口只放在测试里，不留在正式代码中 +alwaysApply: true +--- + +# 测试代码不进正式实现 + +正式路径不保留只为测试、benchmark 或对照方便而存在的函数、参数和分支。 + +需要这类行为时，在测试或 benchmark 自己实现，并调用正式接口。正式接口的参数对应真实推理会传入的数据，不为测试场景增加可选参数。 From 1567491dffcf78efb98cff4073e8185ab406fb6e Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 06:00:57 +0000 Subject: [PATCH 49/54] refactor: name replay kernel arguments and pass tile configs explicitly. Decode launches use the kernel parameter names, and a pinned tile config is an argument instead of a field on the test state. Co-authored-by: Cursor --- .../attention/linear/replay_gdn_decode.py | 74 +++++++------- .../attention/linear/replay_kda_decode.py | 14 ++- test/benchmark/kernels/benchmark_kda_ssm.py | 13 +-- .../linear_att/replay_history.py | 37 +++++-- .../linear_att/test_acceptance_fusion.py | 19 +++- .../linear_att/test_replayssm.py | 98 ++++++++++++++----- .../linear_att/test_replayssm_execution.py | 94 ++++++++++++++---- .../common/test_linear_state_contract.py | 8 +- 8 files changed, 250 insertions(+), 107 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 46a5b09795..5e32cdea8e 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -145,43 +145,43 @@ def replay( precompute_projections = state.compute_mode == "precompute" and bv >= 16 padded_verify_width = triton.next_power_of_2(state.verify_width) _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( - q, - k, - v, - a, - b, - a_log, - bias, - state.ssm_state[linear_layer_idx], - state.keys[linear_layer_idx], - state.deltas[linear_layer_idx], - state.gates[linear_layer_idx], - state.raw_keys[linear_layer_idx], - state.raw_values[linear_layer_idx], - state.betas[linear_layer_idx], - reqs, - state.cursors, - state.tile_sync_mark, - cu_seqlens, - out, - q.stride(0), - k.stride(0), - v.stride(0), - a.stride(0), - b.stride(0), - q.shape[-2], - hv, - state.num_key_heads, - kd, - vd, - state.capacity, - state.hold_req_id, - triton.next_power_of_2(kd), - bv, - precompute_projections, - padded_verify_width, - False, - state.gate_lower_bound, + Q=q, + StepK=k, + StepV=v, + A=a, + B=b, + Alog=a_log, + Bias=bias, + State=state.ssm_state[linear_layer_idx], + Keys=state.keys[linear_layer_idx], + Deltas=state.deltas[linear_layer_idx], + Gates=state.gates[linear_layer_idx], + RawKeys=state.raw_keys[linear_layer_idx], + RawValues=state.raw_values[linear_layer_idx], + Betas=state.betas[linear_layer_idx], + Reqs=reqs, + Cursors=state.cursors, + SyncMark=state.tile_sync_mark, + Cu=cu_seqlens, + Out=out, + SQ=q.stride(0), + SK=k.stride(0), + SV=v.stride(0), + SA=a.stride(0), + SB=b.stride(0), + H=q.shape[-2], + HV=hv, + KH=state.num_key_heads, + K=kd, + V=vd, + L=state.capacity, + HOLD=state.hold_req_id, + BK=triton.next_power_of_2(kd), + BV=bv, + PRECOMPUTE_PROJECTIONS=precompute_projections, + BW=padded_verify_width, + KDA=False, + LOWER=state.gate_lower_bound, BQK=max(16, 2 * padded_verify_width), BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), NV=triton.cdiv(vd, bv), diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index eedbec91b7..af2ed9702b 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -44,12 +44,18 @@ def get_static_key( } -def get_run_key(q, cu_seqlens) -> tuple[int, int]: +def get_run_key(q: torch.Tensor, cu_seqlens: torch.Tensor) -> tuple[int, int]: """调优键是 token 数和请求数。q 是 [tokens, heads, dim]。""" return q.shape[0], cu_seqlens.numel() - 1 -def _check_replay_inputs(q, k, v, reqs, cu_seqlens) -> None: +def _check_replay_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> None: assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" @@ -81,7 +87,7 @@ def rebuild_inputs( ) -> tuple[tuple, dict[str, object]]: """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" del state_cursors, reqs, run_config - tokens, run_count = get_run_key(q, cu_seqlens) + tokens, run_count = get_run_key(q=q, cu_seqlens=cu_seqlens) # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state_verify_width)) slots = batch + 1 @@ -172,7 +178,7 @@ def replay( 宽度为 1 时它是 0, 1, 2, ...。 """ - _check_replay_inputs(q, k, v, reqs, cu_seqlens) + _check_replay_inputs(q=q, k=k, v=v, reqs=reqs, cu_seqlens=cu_seqlens) if run_config is None: run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} _, hv, kd, vd = state_ssm.shape diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index bd12f94321..c60038af40 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -53,6 +53,7 @@ def workload(args, batch, width, dtype, mode, acceptance): ) native_index = 0 if acceptance == "partial" else width - 1 + pinned_config = {"value": None} def cycle(): for _ in range(args.rounds): @@ -94,7 +95,7 @@ def cycle(): state_verify_width=ssm_backend.verify_width, state_compute_mode=ssm_backend.compute_mode, state_gate_lower_bound=ssm_backend.gate_lower_bound, - run_config=ssm_backend.run_config, + run_config=pinned_config["value"], ) if ssm_backend is not None: accept_replay_window(ssm_backend, reqs, accepted, replay_cu) @@ -104,7 +105,7 @@ def cycle(): else: snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) - return ssm_backend, cycle, bytes_used + return ssm_backend, cycle, bytes_used, pinned_config def main(): @@ -133,7 +134,7 @@ def main(): for width in args.width: for acceptance in ["full", "partial"] if width > 1 else ["full"]: for mode in ["native", "inline", "precompute"] if width > 1 else ["native", "inline"]: - ssm_backend, cycle, size = workload( + ssm_backend, cycle, size, pinned_config = workload( args, batch, width, getattr(torch, dtype_name), mode, acceptance ) candidates = [] @@ -141,11 +142,11 @@ def main(): for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: if mode == "precompute" and config["BV"] < 16: continue - ssm_backend.run_config = config + pinned_config["value"] = 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"] + pinned_config["value"] = min(candidates, key=lambda r: r["ms"])["config"] cycle() samples = [ triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) @@ -160,7 +161,7 @@ def main(): state_history_bytes=size, ms_per_round=statistics.median(samples), samples_ms=samples, - config=dict(ssm_backend.run_config) if ssm_backend and ssm_backend.run_config else None, + config=dict(pinned_config["value"]) if pinned_config["value"] else None, candidates=candidates, ) rows.append(row) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index e6fdd1fda3..6714729311 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -17,7 +17,6 @@ def attach_replay_history( num_key_heads: int | None = None, gate_lower_bound: float = -5.0, compute_mode: str = "inline", - run_config: dict[str, int] | None = None, *, kda: bool = False, ): @@ -49,7 +48,6 @@ def attach_replay_history( owner.capacity = capacity owner.verify_width = verify_width owner.compute_mode = compute_mode - owner.run_config = run_config layers, slots, hv, k, v = ssm.shape owner.num_key_heads = hv if num_key_heads is None else num_key_heads assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 @@ -90,15 +88,25 @@ def _trim_rejected_index( tl.store(cursors + r, cursor - (window - accepted), valid) -def replay_layer(state, *args, **kwargs): +def replay_layer( + state: ReplayGDNState | ReplayKDAState, + layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + run_config: dict[str, int] | None = None, +): """按测试 state 的类型调用对应的 replay decode。""" - if "run_config" not in kwargs and state.run_config is not None: - kwargs["run_config"] = state.run_config if state.kda: from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay - layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens = args return replay( state_ssm=state.ssm_state[layer_idx], state_keys=state.keys[layer_idx], @@ -121,11 +129,24 @@ def replay_layer(state, *args, **kwargs): state_verify_width=state.verify_width, state_compute_mode=state.compute_mode, state_gate_lower_bound=state.gate_lower_bound, - run_config=kwargs.get("run_config"), + run_config=run_config, ) from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay - return replay(state, *args, **kwargs) + return replay( + state=state, + linear_layer_idx=layer_idx, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=a_log, + bias=bias, + reqs=reqs, + cu_seqlens=cu_seqlens, + run_config=run_config, + ) def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens=None) -> None: 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 index ddb02125ca..7fbb63df5d 100644 --- 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 @@ -173,7 +173,6 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput torch.bfloat16, num_key_heads=2, compute_mode=compute_mode, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) for _ in range(2) ] @@ -198,7 +197,23 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput def cycle(i): ssm_backend = ssm_backends[i] - outputs = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] + outputs = [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + b, + alog, + alog, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(2) + ] (accept_unfused if i == 0 else accept_fused)(managers[i], args) return outputs 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 index d77fc5b5d8..2d1dec3e41 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -60,7 +60,6 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state num_key_heads=1, kda=kda, compute_mode="precompute", - run_config={"BV": bv, "num_warps": 1, "num_stages": 3}, ) reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") @@ -83,7 +82,20 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state # an overflowed history correction by reconstructing from raw inputs. assert ssm_backend.cursors[0].item() == 2 * round_idx assert not state.any() - actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu) + actual = replay_layer( + ssm_backend, + 0, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": bv, "num_warps": 1, "num_stages": 3}, + ) expected = [] for token in range(2): output, expected_state = reference( @@ -110,7 +122,6 @@ def test_flush_export_updates_checkpoint_and_clears_history(dtype): capacity=4, verify_width=1, activation_dtype=torch.bfloat16, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) backend.raw_keys.normal_(0, 0.1) backend.raw_values.normal_(0, 0.1) @@ -162,7 +173,6 @@ def test_replay_acceptance_flush_and_graph( kda=kda, gate_lower_bound=gate_lower_bound, compute_mode=compute_mode, - run_config=run_config, activation_dtype=activation_dtype, ) expected = state.float().clone() @@ -186,7 +196,10 @@ def test_replay_acceptance_flush_and_graph( accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): - out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=run_config) + for layer in range(2) + ] accept_replay_window(ssm_backend, reqs, accepted, cu) return out @@ -324,7 +337,6 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d dtype, num_key_heads=2, compute_mode=mode, - run_config=config, ) for mode in ("inline", "precompute") ] @@ -358,7 +370,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d for ssm_backend in ssm_backends: ssm_backend.fold_accepted_prefix(overflow_reqs) for ssm_backend in ssm_backends: - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu, run_config=config)) accept_replay_window(ssm_backend, reqs, accepted, cu) if dtype == torch.float32: torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) @@ -420,7 +432,6 @@ def check(capacity, width, steps): torch.bfloat16, num_key_heads=heads, compute_mode=mode, - run_config=config, ) for mode in ("inline", "precompute") ] @@ -433,7 +444,9 @@ def check(capacity, width, steps): v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 b = torch.randn_like(a) - outputs = [replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) for backend in backends] + outputs = [ + replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for backend in backends + ] # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) for backend in backends: @@ -458,8 +471,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): hv, kh, dim = 2, 2, 32 state = torch.randn(2, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 config = {"BV": 32, "num_warps": 1, "num_stages": 1} - backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) - twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda, run_config=config) + backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda) + twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda) assert backend.gates.shape[-1] == (dim if kda else capacity) reqs = torch.tensor([0], device="cuda", dtype=torch.int32) cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") @@ -477,8 +490,8 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): assert int(backend.cursors[0]) == capacity outs, twin_outs = [], [] for layer in range(2): - outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu)) - twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu)) + outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) + twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) if step == capacity: assert int(backend.cursors[0]) == capacity for actual, expected in zip(outs, twin_outs): @@ -495,9 +508,7 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): state = torch.randn(1, 3, 2, 32, 32, device="cuda") - backend = ReplaySSMExecutor( - state, 8, 4, torch.float32, num_key_heads=2, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} - ) + backend = ReplaySSMExecutor(state, 8, 4, torch.float32, num_key_heads=2, kda=kda) backend.cursors[0] = 3 reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) @@ -507,7 +518,20 @@ def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") b = torch.zeros(1, 2, device="cuda") bias = torch.zeros(2 * 32 if kda else 2, device="cuda") - replay_layer(backend, 0, q, k, v, a, b, bias, bias, reqs, cu) + replay_layer( + backend, + 0, + q, + k, + v, + a, + b, + bias, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) accept_replay_window(backend, reqs, accepted, cu) assert backend.cursors.tolist() == [3, 0, 0] @@ -528,7 +552,6 @@ def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_ num_key_heads=2, kda=kda, compute_mode=compute_mode, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ) expected = state.float().clone() pending = [0, 0] @@ -544,7 +567,23 @@ def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_ bias = torch.zeros(4 * 128 if kda else 4, device="cuda") def forward(): - return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] + return [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + for layer in range(2) + ] def step(): positions = ssm_backend.cursors.index_select(0, reqs.long()) @@ -751,9 +790,7 @@ def test_kda_tuning_candidates_match_reference(config, compute_mode): def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): torch.manual_seed(72) initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 - backend = ReplaySSMExecutor( - initial.clone(), 8, 4, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} - ) + backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) reqs = torch.cat( ( torch.arange(live, device="cuda", dtype=torch.int32), @@ -787,6 +824,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa bias, reqs, cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ).reshape(live, query_width, 1, 32) for token in range(query_width): output, expected = reference( @@ -825,7 +863,6 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): num_key_heads=kh, kda=kda, compute_mode=compute_mode, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ) reqs = torch.tensor([0], device="cuda", dtype=torch.int32) indices = torch.zeros(2, device="cuda", dtype=torch.int32) @@ -838,7 +875,20 @@ def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): v = torch.randn(2, hv, dim, device="cuda", dtype=dtype) a = torch.full((2, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) b = torch.randn(2, hv, device="cuda", dtype=dtype) - replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu) + replay_layer( + backend, + 0, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) lengths = torch.tensor([2], device="cuda", dtype=torch.int32) advance_cursor(backend, reqs, cu) commit_ssm_prefix(indices, reqs, lengths, lengths, backend.cursors) 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 index ef4a1c618c..02797b61b7 100644 --- 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 @@ -36,12 +36,7 @@ def record(*args, run=kernel.run, name=name, **kwargs): torch.manual_seed(137) initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 - eager, captured = [ - ReplaySSMExecutor( - initial.clone(), 8, 1, torch.float32, kda=kda, run_config={"BV": 32, "num_warps": 1, "num_stages": 3} - ) - for _ in range(2) - ] + eager, captured = [ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda) for _ in range(2)] reqs = torch.arange(512, device="cuda", dtype=torch.int32) q = torch.randn(512, 1, 32, device="cuda") * 0.2 k, v = torch.randn_like(q), torch.randn_like(q) @@ -52,7 +47,20 @@ def record(*args, run=kernel.run, name=name, **kwargs): def cycle(backend, ids): count = ids.numel() cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) - out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids, cu) + out = replay_layer( + backend, + 0, + q[:count], + k[:count], + v[:count], + a[:count], + beta[:count], + alog, + bias, + ids, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) accept_replay_window(backend, ids, cu_seqlens=cu) return out @@ -106,7 +114,6 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle num_key_heads=2, kda=kda, compute_mode="precompute", - run_config=config, ) for cls in [ReplaySSMExecutor, ReplaySSMExecutor] ] @@ -155,7 +162,10 @@ def metadata(ssm_backend): def forward_accept(index, cu, reqs): ssm_backend = ssm_backends[index] - out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(3) + ] if index == 0: accept_replay_window(ssm_backend, reqs, accepted, cu) else: @@ -247,7 +257,6 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, dtype, num_key_heads=2, kda=kda, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) for _ in range(2) ] @@ -265,7 +274,23 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) def forward(ssm_backend): - return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] + return [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(3) + ] def warm(ssm_backend): out = forward(ssm_backend) @@ -336,8 +361,7 @@ def test_cross_layer_fold_with_state_offset_above_int32(kda): sparse.copy_(dense) config = {"BV": 32, "num_warps": 1, "num_stages": 1} ssm_backends = [ - ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda, run_config=config) - for state in [dense, sparse] + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda) for state in [dense, sparse] ] q = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) @@ -354,7 +378,10 @@ def step(ssm_backend, accepted): if overflow: ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) pos = ssm_backend.cursors.index_select(0, reqs.long()) - out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(48)] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(48) + ] accept_replay_window(ssm_backend, reqs, accepted, cu) return pos, out @@ -383,8 +410,7 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 config = {"BV": 32, "num_warps": 1, "num_stages": 1} ssm_backends = [ - ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) - for _ in range(2) + ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute") for _ in range(2) ] reqs = torch.arange(batch, device="cuda", dtype=torch.int32) padded = torch.cat((reqs, reqs.new_full((13 - batch,), 13))) @@ -402,7 +428,10 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): def step(): - out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets) for layer in range(2)] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets, run_config=config) + for layer in range(2) + ] accept_replay_window(ssm_backend, ids, accepted, offsets) return out @@ -446,7 +475,6 @@ def test_full_prefix_does_not_address_past_capacity(kda): 1, torch.bfloat16, kda=kda, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) @@ -460,13 +488,39 @@ def test_full_prefix_does_not_address_past_capacity(kda): for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): getattr(backend, name).zero_() backend.cursors[reqs.long()] = capacity + 2 - replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs, cu) + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) backend.fold_accepted_prefix(reqs) torch.cuda.synchronize() assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) assert torch.isfinite(backend.ssm_state).all() for _ in range(12): - replay_layer(backend, 0, q, k, v, a, beta, alog, bias, reqs, cu) + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) accept_replay_window(backend, reqs) overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] if overflow.numel(): diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 33c4b6be23..fbfc4a43a9 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -173,12 +173,8 @@ def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, modul torch.arange(2, dtype=torch.int32), run_config=run_config, ) - if cls is ReplayKDAState: - assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute - assert captured.kwargs["KDA"] is True - else: - assert captured.args[33] is precompute - assert captured.args[35] is False + assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute + assert captured.kwargs["KDA"] is (cls is ReplayKDAState) @pytest.mark.parametrize("cls", [GDNState, KDAState]) From 7bcedeffcd3174f6275afe7e4ac60dd0a0a537e8 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 06:03:30 +0000 Subject: [PATCH 50/54] refactor: move SSM config validation to the end of startup. The helper is only called after the module loads, so it sits with the other startup checks instead of above the launch path. Co-authored-by: Cursor --- lightllm/server/api_start.py | 102 +++++++++++++++++------------------ 1 file changed, 51 insertions(+), 51 deletions(-) diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 95379e5447..a4c990c2be 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -43,57 +43,6 @@ 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", "replay"), "Invalid ssm_state_mode" - config = args.ssm_replay_config - assert isinstance(config, dict), "ssm_replay_config must be a JSON object" - assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" - if args.ssm_state_mode == "replay": - defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} - unknown = config.keys() - defaults.keys() - assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" - config = {**defaults, **config} - assert type(config["cache_len"]) is int and config["cache_len"] in ( - 4, - 8, - 16, - 32, - 64, - ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" - assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" - args.ssm_replay_config = config - logger.info(f"SSM mode: replay, config: {config}") - else: - logger.info(f"SSM mode: {args.ssm_state_mode}") - from lightllm.utils.config_utils import get_model_type - - model_type = get_model_type(args.model_dir) - is_gdn_or_kda = model_type in ( - "glm5_next", - "glm5_next_text", - "qwen3_next", - "qwen3_5", - "qwen3_5_moe", - "qwen3_5_text", - "qwen3_5_moe_text", - "sensenova_loop", - ) - if is_gdn_or_kda: - assert args.mtp_step <= 15, ( - "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" - ) - if args.ssm_state_mode == "native": - return - assert args.linear_att_ssm_data_type in ( - "float32", - "bfloat16", - ), "Non-native SSM state modes require FP32 or BF16 state" - assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" - assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" - capacity = config["cache_len"] - assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" - - def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) _set_envs_and_config(args) @@ -797,3 +746,54 @@ def _auto_set_prefill_queue_strategy(args: StartArgs) -> None: if args.prefill_queue_strategy == "auto": args.prefill_queue_strategy = "hrrn" if args.run_mode == "prefill" else "default" logger.info(f"run_mode={args.run_mode}, auto set prefill_queue_strategy to {args.prefill_queue_strategy}") + + +def _validate_and_normalize_ssm_config(args: StartArgs): + assert args.ssm_state_mode in ("native", "replay"), "Invalid ssm_state_mode" + config = args.ssm_replay_config + assert isinstance(config, dict), "ssm_replay_config must be a JSON object" + assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" + if args.ssm_state_mode == "replay": + defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} + unknown = config.keys() - defaults.keys() + assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" + config = {**defaults, **config} + assert type(config["cache_len"]) is int and config["cache_len"] in ( + 4, + 8, + 16, + 32, + 64, + ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" + assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" + args.ssm_replay_config = config + logger.info(f"SSM mode: replay, config: {config}") + else: + logger.info(f"SSM mode: {args.ssm_state_mode}") + from lightllm.utils.config_utils import get_model_type + + model_type = get_model_type(args.model_dir) + is_gdn_or_kda = model_type in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "sensenova_loop", + ) + if is_gdn_or_kda: + assert args.mtp_step <= 15, ( + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" + ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" + capacity = config["cache_len"] + assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" From 5a441ba3e8cd5d12585101405930c47f01a83e71 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 08:16:18 +0000 Subject: [PATCH 51/54] fix: ignore reused replay history and store deltas in fp32 A restored slot only receives the canonical Conv and SSM, so its cursor must be cleared or the previous request's history is folded in. Deltas can exceed the fp16 range, so keep them in fp32 while normalized keys stay fp16 for bf16 activations. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 21 +++--- .../bounded_att/kda/replay_kda_state.py | 21 +++--- .../linear_att/replay_history.py | 4 +- .../linear_att/test_replayssm.py | 18 +++-- unit_tests/common/test_linear_state.py | 75 +++++++++++++++++++ 5 files changed, 113 insertions(+), 26 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index ab2bfaf490..12a948cf89 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -103,13 +103,13 @@ def __init__( self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) - # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 + # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # GDN 的 g 是每个头一个标量。 reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, self.num_key_heads, capacity, k) self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) - self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=device) self.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) @@ -206,7 +206,7 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - """从页恢复基准 Conv/SSM。prefill 时 cursor 仍是 0。""" + """从页恢复基准 Conv/SSM。使目标槽的旧 history 失效。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) @@ -217,8 +217,9 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_enabled: - for req_idx in req_indices: + for req_idx in req_indices: + self.cursors[req_idx].zero_() + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state( @@ -333,7 +334,7 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayGDNState", ) -> None: - """复制 prefill 基准 Conv/SSM 到目标单槽。prefill 时目标 cursor 仍是 0。""" + """复制 prefill 基准 Conv/SSM 到目标单槽。使目标槽的旧 history 失效。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" assert ( @@ -351,8 +352,9 @@ def copy_runtime_state_from_dp( ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_enabled: - for dst_req in dst_req_indices: + for dst_req in dst_req_indices: + self.cursors[dst_req].zero_() + if self.mtp_enabled: self.mtp_state_index[dst_req].fill_(0) # ======================================================================== @@ -397,7 +399,7 @@ def read_from_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """读入基准 Conv/SSM。prefill 时 cursor 仍是 0。""" + """读入基准 Conv/SSM。使目标槽的旧 history 失效。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._tensor_fields( @@ -412,6 +414,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) + self.cursors[req_idx].zero_() if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index cbb068b164..fec80ea755 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -103,13 +103,13 @@ def __init__( self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) - # 输出重建用归一化 k 和 d。bf16 激活用 fp16，尾数更多。 + # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # KDA 的 g 按 key 通道存。 reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, self.num_key_heads, capacity, k) self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) - self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=device) self.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=device) self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) @@ -206,7 +206,7 @@ def restore_bounded_att_runtime_from_pages( page_kind, page_slot_ids: list[int], ) -> None: - """从页恢复基准 Conv/SSM。prefill 时 cursor 仍是 0。""" + """从页恢复基准 Conv/SSM。使目标槽的旧 history 失效。""" assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) @@ -217,8 +217,9 @@ def restore_bounded_att_runtime_from_pages( ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_enabled: - for req_idx in req_indices: + for req_idx in req_indices: + self.cursors[req_idx].zero_() + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state( @@ -333,7 +334,7 @@ def copy_runtime_state_from_dp( dst_req_indices: list[int], remote_dp_bounded_state: "ReplayKDAState", ) -> None: - """复制 prefill 基准 Conv/SSM 到目标单槽。prefill 时目标 cursor 仍是 0。""" + """复制 prefill 基准 Conv/SSM 到目标单槽。使目标槽的旧 history 失效。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" assert ( @@ -351,8 +352,9 @@ def copy_runtime_state_from_dp( ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_enabled: - for dst_req in dst_req_indices: + for dst_req in dst_req_indices: + self.cursors[dst_req].zero_() + if self.mtp_enabled: self.mtp_state_index[dst_req].fill_(0) # ======================================================================== @@ -397,7 +399,7 @@ def read_from_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """读入基准 Conv/SSM。prefill 时 cursor 仍是 0。""" + """读入基准 Conv/SSM。使目标槽的旧 history 失效。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._tensor_fields( @@ -412,6 +414,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) + self.cursors[req_idx].zero_() if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 6714729311..ae6037853e 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -55,12 +55,12 @@ def attach_replay_history( owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) owner.tile_sync_mark = torch.zeros((slots * (owner.mtp_step + 6), hv), dtype=torch.int32, device=ssm.device) - # 重建用 k/d：bf16 激活用 fp16。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 + # 重建用归一化 k：bf16 激活用 fp16；d 使用 fp32。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) gate_shape = (layers, slots, hv, capacity, k) if kda else (layers, slots, hv, capacity) owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) - owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) owner.gates = torch.empty(gate_shape, dtype=torch.float32, device=ssm.device) owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) 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 index 2d1dec3e41..b6a918a5b2 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -50,13 +50,14 @@ def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5. @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): +@pytest.mark.parametrize("activation_dtype", [torch.float16, torch.bfloat16]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda, activation_dtype): state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") ssm_backend = ReplaySSMExecutor( state, capacity, 2, - torch.float16, + activation_dtype, num_key_heads=1, kda=kda, compute_mode="precompute", @@ -64,15 +65,15 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") accepted = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - q = torch.zeros(4, 1, 128, dtype=torch.float16, device="cuda") + q = torch.zeros(4, 1, 128, dtype=activation_dtype, device="cuda") q[:2, 0, 0] = 1 k = torch.zeros_like(q) k[0, 0, 0] = -0.25 k[1, 0, 0] = 1 v = torch.zeros_like(q) v[:2, 0, 0] = 60000 - a = torch.full((4, 128 if kda else 1), -20, dtype=torch.float16, device="cuda") - b = torch.full((4, 1), 20, dtype=torch.float16, device="cuda") + a = torch.full((4, 128 if kda else 1), -20, dtype=activation_dtype, device="cuda") + b = torch.full((4, 1), 20, dtype=activation_dtype, device="cuda") alog = torch.zeros(1, device="cuda") bias = torch.zeros(128 if kda else 1, device="cuda") expected_state = state[0, 0].float().clone() @@ -104,7 +105,12 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state expected.append(output) assert torch.isfinite(expected_state).all() assert torch.isfinite(actual).all() - torch.testing.assert_close(actual[:2].float(), torch.stack(expected), atol=4, rtol=1e-3) + torch.testing.assert_close( + actual[:2].float(), + torch.stack(expected), + atol=32 if activation_dtype == torch.bfloat16 else 4, + rtol=4e-3 if activation_dtype == torch.bfloat16 else 1e-3, + ) accept_replay_window(ssm_backend, reqs, accepted, cu) # The correction is larger than FP16 can represent even though the # recurrent state, inputs, and attention outputs are all finite. diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index bb2314a025..e538002a17 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -79,6 +79,9 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), lambda: state.read_from_pd_page(block, dst, 0, 8, layout), + lambda: state.copy_runtime_state_from_dp( + src_req_indices=[src], dst_req_indices=[dst], remote_dp_bounded_state=state + ), ): if state.mtp_enabled: state.mtp_state_index[dst] = step @@ -96,6 +99,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( assert state.mtp_state_index[dst].item() == 0 if mode == "replay": assert state.cursors[dst].item() == 0 + assert state.cursors[src].item() == 1 state.reset_bounded_att_state([dst]) assert not state.conv_state[:, dst].any() if mode == "replay": @@ -257,3 +261,74 @@ def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(mo other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) assert torch.equal(block, exported) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("step", [0, 2]) +@pytest.mark.parametrize("restore_kind", ["checkpoint", "pd", "dp"]) +def test_replay_restored_slot_decodes_like_clean_slot(runtime_start_args, cls, step, restore_kind): + from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import replay_layer + + torch.manual_seed(1690) + state, _, layout = make_state(cls=cls, mode="replay", step=step) + state.initialize_page_pools(big_page_capacity=1) + state.ssm_state.normal_(mean=0, std=0.02) + state.conv_state.normal_(mean=0, std=0.1) + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[8], page_kind="big", page_slot_ids=[0]) + state.write_to_pd_page(req_idx=0, pd_page_buffers=block, pd_page_index=0, state_token_len=8, layout=layout) + hv, kd, vd = state.ssm_state.shape[-3:] + kh = state.num_key_heads + width = step + 1 + q = torch.randn(width, kh, kd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + k = torch.randn_like(q[:width]).repeat(2, 1, 1) + v = torch.randn(width, hv, vd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + a = torch.full((2 * width, hv * kd if state.kda else hv), -3.0, device="cuda") + b = torch.zeros(2 * width, hv, device="cuda") + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(a.shape[-1], device="cuda") + reqs = torch.tensor([1, 2], dtype=torch.int32, device="cuda") + cu = torch.arange(3, dtype=torch.int32, device="cuda") * width + + def decode(): + outputs = [] + for layer in state.config.logical_layer_ids: + outputs.append( + replay_layer( + state=state, + layer_idx=layer, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=alog, + bias=bias, + reqs=reqs, + cu_seqlens=cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + ) + state.finish_decode_token(logical_layer_id=layer, reqs=reqs, cu_seqlens=cu) + return outputs + + decode() + assert state.cursors[1].item() == width + state.reset_bounded_att_state(req_indices=[2]) + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + if restore_kind == "checkpoint": + state.restore_bounded_att_runtime_from_pages( + req_indices=[1], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + elif restore_kind == "pd": + state.read_from_pd_page(pd_page_buffers=block, req_idx=1, pd_page_index=0, state_token_len=8, layout=layout) + else: + state.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=state) + assert state.cursors[1].item() == 0 + for _ in range(3): + for output in decode(): + torch.testing.assert_close(output[:width], output[width:], rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], state.ssm_state[:, 2], rtol=0, atol=0) From 6fb01a0935b8dd6be916d24fe9cd08520cb07053 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 08:43:51 +0000 Subject: [PATCH 52/54] fix: wait for every value head sharing a key before replay overwrites it Grouped GDN configs read one RawKeys buffer from several value heads. The fold barrier now counts that whole key group, and the sync mark is shaped to the key heads. Co-authored-by: Cursor --- .../bounded_att/gdn/replay_gdn_state.py | 6 ++-- .../bounded_att/kda/replay_kda_state.py | 6 ++-- .../attention/linear/replay_gdn_decode.py | 11 ++++++- .../attention/linear/replay_kda_decode.py | 13 ++++++-- .../linear_att/replayssm_decode.py | 33 ++++++++++++++----- .../linear_att/replay_history.py | 4 ++- 6 files changed, 57 insertions(+), 16 deletions(-) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 12a948cf89..137ce8ad32 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -101,8 +101,10 @@ def __init__( self.hold_req_id = slots - 1 # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) - # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 - self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) + # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros( + (slots * (mtp_step + 6), self.num_key_heads), dtype=torch.int32, device=device + ) # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # GDN 的 g 是每个头一个标量。 diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index fec80ea755..2c73e97d6f 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -101,8 +101,10 @@ def __init__( self.hold_req_id = slots - 1 # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) - # 按 launch 行号同步 V tile。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 - self.tile_sync_mark = torch.zeros((slots * (mtp_step + 6), hv), dtype=torch.int32, device=device) + # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros( + (slots * (mtp_step + 6), self.num_key_heads), dtype=torch.int32, device=device + ) # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 # KDA 的 g 按 key 通道存。 diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index 5e32cdea8e..cef9fe745f 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -11,7 +11,10 @@ import triton from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( + _decode_replay, + check_replay_sync_mark, +) from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune @@ -144,6 +147,12 @@ def replay( # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 precompute_projections = state.compute_mode == "precompute" and bv >= 16 padded_verify_width = triton.next_power_of_2(state.verify_width) + check_replay_sync_mark( + sync_mark=state.tile_sync_mark, + reqs=reqs, + value_heads=hv, + key_heads=state.num_key_heads, + ) _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( Q=q, StepK=k, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index af2ed9702b..1f700b65b0 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -10,7 +10,10 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( + _decode_replay, + check_replay_sync_mark, +) from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune @@ -116,7 +119,7 @@ def rebuild_inputs( state_raw_values=torch.empty(slots, hv, capacity, vd, device=device, dtype=state_raw_values.dtype), state_betas=torch.empty(slots, hv, capacity, device=device, dtype=state_betas.dtype), state_cursors=torch.zeros(slots, dtype=torch.int32, device=device), - state_tile_sync_mark=torch.zeros((slots * rows_per_slot, hv), dtype=torch.int32, device=device), + state_tile_sync_mark=torch.zeros((slots * rows_per_slot, key_heads), dtype=torch.int32, device=device), q=copied[0], k=copied[1], v=copied[2], @@ -190,6 +193,12 @@ def replay( # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 precompute_projections = state_compute_mode == "precompute" and bv >= 16 padded_verify_width = triton.next_power_of_2(state_verify_width) + check_replay_sync_mark( + sync_mark=state_tile_sync_mark, + reqs=reqs, + value_heads=hv, + key_heads=key_heads, + ) _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( Q=q, StepK=k, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index 8f16764d54..bdc9b02d81 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -6,21 +6,38 @@ state. The output is either a precomputed GDN or KDA window, or a token walk. """ +import torch import triton import triton.language as tl from .replayssm_delta import _apply_accepted_prefix +def check_replay_sync_mark( + sync_mark: torch.Tensor, + reqs: torch.Tensor, + value_heads: int, + key_heads: int, +) -> None: + """确认 SyncMark 按 [launch 行, key head] 连续存放，并且盖住这次启动的每一行。""" + + assert key_heads > 0 and value_heads % key_heads == 0, "value heads must be a multiple of key heads" + assert sync_mark.dtype == torch.int32 and sync_mark.dim() == 2, "sync mark must be a 2D int32 tensor" + assert sync_mark.shape[1] == key_heads, "sync mark columns must equal key heads" + assert sync_mark.shape[0] >= reqs.numel(), "sync mark rows must cover every launch row" + assert sync_mark.stride(1) == 1 and sync_mark.stride(0) == key_heads, "sync mark must be contiguous by key head" + + @triton.jit -def _arrive_value_tiles(SyncMark, seq, hv, HV: tl.constexpr, NV: tl.constexpr): - # 同一个 launch 行、同一个 value head 的 V/BV 个 tile 都读完旧前缀后，才允许写新窗口。 +def _arrive_value_tiles(SyncMark, seq, hv, HV: tl.constexpr, KH: tl.constexpr, NV: tl.constexpr): + # 同一 launch 行、同一 key head 的全部 value head 和 V tile 都读完旧 RawKeys 后，才允许写新窗口。 # 行号不是请求槽。HOLD 会在一批里重复出现，用请求槽做下标会撞到同一格。 - # 最后到达的 tile 把计数写回 0，这个 0 就是放行信号。其余 tile 已经加过 1， - # 之后再看见 0 只可能是这次复原。 - idx = seq * HV + hv + # SyncMark 每行长度是 KH。最后到达的 tile 把计数写回 0，这个 0 就是放行信号。 + group: tl.constexpr = HV // KH + key_head = hv // group + idx = seq * KH + key_head arrived = tl.atomic_add(SyncMark + idx, 1, sem="acq_rel", scope="gpu") - if arrived == NV - 1: + if arrived == NV * group - 1: tl.atomic_xchg(SyncMark + idx, 0, sem="release", scope="gpu") else: seen = arrived + 1 @@ -117,8 +134,8 @@ def _decode_replay( state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) n = 0 - # 新 token 会从槽 0 覆盖刚读过的 RawKeys、Gates、Betas。等这一组 V tile 都读完再往下走。 - _arrive_value_tiles(SyncMark, seq, hv, HV, NV) + # 新 token 会从槽 0 覆盖刚读过的 RawKeys。等同一 key head 的全部 value tile 都读完再往下走。 + _arrive_value_tiles(SyncMark, seq, hv, HV, KH, NV) if PRECOMPUTE_PROJECTIONS: # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index ae6037853e..956bf5e173 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -54,7 +54,9 @@ def attach_replay_history( # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度。 owner.hold_req_id = slots - 1 owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) - owner.tile_sync_mark = torch.zeros((slots * (owner.mtp_step + 6), hv), dtype=torch.int32, device=ssm.device) + owner.tile_sync_mark = torch.zeros( + (slots * (owner.mtp_step + 6), owner.num_key_heads), dtype=torch.int32, device=ssm.device + ) # 重建用归一化 k：bf16 激活用 fp16；d 使用 fp32。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 key_shape = (layers, slots, owner.num_key_heads, capacity, k) From 59adc054c637bf1100c3ec0b5245958824cf6314 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 12 Nov 2026 08:45:47 +0000 Subject: [PATCH 53/54] refactor: launch replay value heads before requests Neighboring programs now walk the V tiles and then the value heads of one request. Those heads sit in adjacent SSM slots and share the key-head fold barrier. Co-authored-by: Cursor --- .../common/basemodel/attention/linear/replay_gdn_decode.py | 2 +- .../common/basemodel/attention/linear/replay_kda_decode.py | 2 +- .../basemodel/triton_kernel/linear_att/replayssm_decode.py | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py index cef9fe745f..ac7af9dbdc 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -153,7 +153,7 @@ def replay( value_heads=hv, key_heads=state.num_key_heads, ) - _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + _decode_replay[(triton.cdiv(vd, bv), hv, reqs.numel())]( Q=q, StepK=k, StepV=v, diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py index 1f700b65b0..3704ba4fd6 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -199,7 +199,7 @@ def replay( value_heads=hv, key_heads=key_heads, ) - _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + _decode_replay[(triton.cdiv(vd, bv), hv, reqs.numel())]( Q=q, StepK=k, StepV=v, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py index bdc9b02d81..f3abbeb020 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -109,7 +109,8 @@ def _decode_replay( 前缀放不下本步 token 时，先在本 kernel 里把已有前缀从原始 k/v 折进 checkpoint， 舍入成 SSM 的 dtype，再从这份状态重放本步。cursor 本 kernel 不修改，留给最后一层返回后更新。 """ - iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) + # grid 是 (V tile, value head, launch 行)。同一 V tile 上相邻 program 对应连续的 value head。 + iv, hv, seq = tl.program_id(0), tl.program_id(1), tl.program_id(2) req = tl.load(Reqs + seq).to(tl.int64) start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) vv = iv * BV + tl.arange(0, BV) From 083ac58801ad544b2ed023bc7433438f1a8efcea Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 17:22:49 +0800 Subject: [PATCH 54/54] test: fix replay padding and restore grouped-head coverage --- .../linear_att/test_replayssm.py | 72 ++++++++++++++++++- .../models/sensenova_loop/test_gdn_state.py | 9 +-- 2 files changed, 76 insertions(+), 5 deletions(-) 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 index b6a918a5b2..786376a362 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -511,6 +511,76 @@ def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): torch.testing.assert_close(backend.ssm_state, twin.ssm_state, atol=state_atol, rtol=0) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("bv", [32, 128]) +@pytest.mark.parametrize("width", [1, 3]) +def test_overflow_preserves_shared_keys_across_heads_and_graph(kda, compute_mode, bv, width): + torch.manual_seed(123) + batch, kh, hv, dim, capacity = 64, 4, 16, 128, 4 + dtype = torch.bfloat16 + initial = torch.randn(2, batch + 1, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + config = {"BV": bv, "num_warps": 4 if bv == 128 else 1, "num_stages": 1} + backend, reference_backend = [ + ReplaySSMExecutor( + initial.clone(), + capacity, + width, + dtype, + num_key_heads=kh, + kda=kda, + compute_mode=compute_mode, + ) + for _ in range(2) + ] + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + reqs[-2:] = batch + windows = (torch.arange(reqs.numel(), device="cuda", dtype=torch.int32) % width) + 1 + windows[0] = 0 + cu = torch.cat((windows.new_zeros(1), windows.cumsum(0).int())) + tokens = int(cu[-1]) + q = torch.randn(tokens, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(tokens, hv, dim, device="cuda", dtype=dtype) + a = torch.full((tokens, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) + b = torch.randn(tokens, hv, device="cuda", dtype=dtype) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + + def step(owner): + outputs = [ + replay_layer(owner, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for layer in range(2) + ] + accept_replay_window(owner, reqs, cu_seqlens=cu) + return outputs + + step(backend) + backend.ssm_state.copy_(initial) + backend.cursors.zero_() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = step(backend) + backend.ssm_state.copy_(initial) + backend.cursors.zero_() + for _ in range(2 * capacity + 1): + for tensor in (q, k, v, b): + tensor.normal_() + lengths = reference_backend.cursors[reqs[:-2].long()] + overflow = reqs[:-2][(windows[:-2] > 0) & (lengths + windows[:-2] > capacity)] + if overflow.numel(): + fold_prefix(state=reference_backend, reqs=overflow, run_config=config) + reference_backend.cursors.index_fill_(0, overflow.long(), 0) + expected = step(reference_backend) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(backend.ssm_state, reference_backend.ssm_state, rtol=0, atol=0) + for left, right in zip(actual, expected): + torch.testing.assert_close(left, right, rtol=torch.finfo(dtype).eps, atol=2e-5) + assert torch.equal(backend.cursors, reference_backend.cursors) + torch.testing.assert_close(backend.ssm_state[:, 0], initial[:, 0], rtol=0, atol=0) + torch.testing.assert_close(backend.ssm_state[:, batch], initial[:, batch], rtol=0, atol=0) + + @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): state = torch.randn(1, 3, 2, 32, 32, device="cuda") @@ -645,7 +715,7 @@ def step(): pending[req] += count assert ssm_backend.cursors.tolist() == pending + [0, 0] for output in outputs: - assert torch.count_nonzero(output[:, sum(widths) :]).item() == 0 + 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 diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index a31c0acaf8..c25576a19a 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -58,7 +58,9 @@ def test_gdn_decode_keeps_logical_layer_state( resources = AttentionStateManager(request_capacity=3) resources.register_state("main", resource) resources.finish_state_registration() - assert resource.get_storage_layer_index(logical_layer) == storage_layer + conv_view, ssm_view = resource.get_layer_view(logical_layer_id=logical_layer) + assert conv_view.data_ptr() == resource.conv_state[storage_layer].data_ptr() + assert ssm_view.data_ptr() == resource.ssm_state[storage_layer].data_ptr() physical_weights = [ SimpleNamespace( layer_num_=i, @@ -94,11 +96,10 @@ def recurrent(*, initial_state, v, **kwargs): "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 + "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent ) monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent." - "mtp_fused_recurrent_gated_delta_rule", + "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", recurrent, )