feat: add ReplaySSM state modes for GDN and KDA

lightllm/common/attention_state/bounded_att/base_state.py

@@ -7,3 +7,7 @@
 
 class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface):
     """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。"""
+
+    def validate_bounded_att_commit(self, context):
+
+        pass

lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py

@@ -13,11 +13,18 @@ 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(
         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,
     )

lightllm/common/attention_state/bounded_att/interface.py

@@ -102,8 +102,9 @@ def commit_bounded_att_state(
         self,
         req_indices: torch.Tensor,
         accepted_steps: torch.Tensor,
+        *,
+        context=None,
     ) -> None:
-        """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。"""
         ...
 
 
@@ -173,10 +174,7 @@ def copy_runtime_state_from_dp(
         dst_req_indices: List[int],
         remote_dp_bounded_state: CoreInterface,
     ) -> None:
-        """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。
 
-        调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。
-        """
         ...
 
 
@@ -205,10 +203,7 @@ def write_to_pd_page(
         state_token_len: int,
         layout: StateTransferLayout,
     ) -> None:
-        """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。
 
-        调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。
-        """
         ...
 
     @abstractmethod

lightllm/common/attention_state/bounded_att/kda/create_kda_state.py

@@ -13,11 +13,18 @@ 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(
         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,
     )

lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py

@@ -179,7 +179,9 @@ def restore_bounded_att_runtime_from_pages(
                 write=False,
             )
 
-    def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None:
+    def commit_bounded_att_state(
+        self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, *, context=None
+    ) -> None:
         """ring 按绝对 token 位置寻址，无需提交额外游标。"""
         pass
 

lightllm/common/attention_state/bounded_att/ssm/__init__.py

@@ -0,0 +1,34 @@
+import torch
+
+
+def create_ssm_state(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound):
+    if mode not in ("native", "compact", "replay"):
+        raise ValueError(f"unknown SSM state mode: {mode}")
+    slots = mtp_step + 1 if mode == "native" else 1
+    state = torch.zeros(
+        (len(config.logical_layer_ids), request_capacity * slots, config.num_v_heads, config.key_dim, config.value_dim),
+        dtype=config.ssm_dtype,
+        device=device,
+    )
+    kwargs = dict(kda=kda, gate_lower_bound=gate_lower_bound)
+    if mode == "native":
+        from .native import NativeSSMState
+
+        return NativeSSMState(state, mtp_step + 1, **kwargs)
+    kwargs.update(num_key_heads=config.num_k_heads)
+    if mode == "compact":
+        from .compact import CompactSSMState
+
+        return CompactSSMState(state, mtp_step + 1, config.conv_dtype, **kwargs)
+    from .replay import ReplaySSMState
+
+    replay_config = replay_config or {}
+    return ReplaySSMState(
+        state,
+        replay_config.get("cache_len", 16),
+        mtp_step + 1,
+        config.conv_dtype,
+        compute_mode=replay_config.get("compute_mode", "inline"),
+        fold_method=replay_config.get("fold_method", "raw"),
+        **kwargs,
+    )

lightllm/common/attention_state/bounded_att/ssm/base.py

@@ -0,0 +1,66 @@
+from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update
+
+
+def update_accepted_indices(reqs, accepted_steps, indices, context, *, replay_cursors=None):
+    if context is None:
+        indices[reqs] = accepted_steps - 1
+    else:
+        linear_att_mtp_state_index_update(
+            req_to_mtp_state_index=indices,
+            b_req_mtp_start_loc=context.b_req_mtp_start_loc,
+            b_req_idx=context.b_req_idx,
+            b_mtp_index=context.b_mtp_index,
+            accepted_index=context.accepted_index,
+            verify_width=context.verify_width,
+            replay_cursors=replay_cursors,
+        )
+
+
+class SSMState:
+    slots_per_request = 1
+
+    def clear_history(self, req):
+        pass
+
+    def prepare_decode(self, reqs, cu_seqlens=None):
+        pass
+
+    def accept_single_token(self, reqs):
+        pass
+
+    def merge_accepted_updates(self, reqs):
+        pass
+
+    def decode_indices(self, reqs, speculative):
+        return reqs
+
+    def get_accepted_state(self, req_idx, accepted_index=None):
+        return self.state[:, req_idx]
+
+    def create_commit(self, run_config):
+        return None
+
+    def validate_commit(self, context, state_name):
+        pass
+
+    def commit(self, reqs, accepted_steps, indices, context, state_name):
+        update_accepted_indices(reqs, accepted_steps, indices, context)
+
+    def forward(
+        self,
+        layer,
+        q,
+        k,
+        v,
+        a,
+        b,
+        a_log,
+        bias,
+        reqs,
+        cu_seqlens=None,
+        run_config=None,
+        *,
+        state_indices=None,
+        accepted_token_counts=None,
+    ):
+        raise NotImplementedError

lightllm/common/attention_state/bounded_att/ssm/compact.py

@@ -0,0 +1,177 @@
+import torch
+import triton
+
+from lightllm.common.ssm_types import SSMRunConfig
+from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import _compact
+from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config
+from .base import SSMState, update_accepted_indices
+
+
+class CompactSSMState(SSMState):
+    def __init__(
+        self,
+        state,
+        verify_width,
+        activation_dtype,
+        *,
+        num_key_heads=None,
+        kda=False,
+        gate_lower_bound=-5.0,
+        run_config=None,
+    ):
+        self.state = state
+        self.kda = kda
+        self.gate_lower_bound = gate_lower_bound
+        self.verify_width = verify_width
+        # Verify and accepted-state reconstruction must share the same layout.
+        self._config_is_fixed = run_config is not None
+        self.batch_configs = {}
+        self.run_config = dict(run_config or {"BV": 8, "num_warps": 1})
+        assert self.run_config["BV"] in (8, 16, 32, 64, 128)
+        assert self.run_config["num_warps"] in (1, 2, 4, 8)
+        layers, slots, hv, k, v = state.shape
+        self.num_key_heads = hv if num_key_heads is None else num_key_heads
+        assert self.num_key_heads > 0 and hv % self.num_key_heads == 0
+        self.hold_req_id = slots - 1
+        shape = (layers, slots, hv, verify_width)
+        self.keys = torch.empty(
+            (layers, slots, self.num_key_heads, verify_width, k), device=state.device, dtype=activation_dtype
+        )
+        self.values = torch.empty((*shape, v), device=state.device, dtype=activation_dtype)
+        self.decays = torch.empty((*shape, k) if kda else shape, device=state.device, dtype=torch.float32)
+        self.betas = torch.empty(shape, device=state.device, dtype=torch.float32)
+
+    def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig:
+        config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens)
+        return SSMRunConfig(**config)
+
+    def forward(
+        self,
+        linear_layer_idx,
+        q,
+        k,
+        v,
+        a,
+        b,
+        a_log,
+        bias,
+        reqs,
+        cu_seqlens=None,
+        run_config=None,
+        *,
+        state_indices=None,
+        accepted_token_counts=None,
+    ):
+        assert cu_seqlens is not None, "compact replay is only used for speculative verify"
+        config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens)
+        if not isinstance(config, SSMRunConfig):
+            config = SSMRunConfig(**config)
+        _, slots, hv, kd, vd = self.state.shape
+        assert self.num_key_heads in (q.shape[-2], hv)
+        out = torch.empty_like(v)
+        # Match native reduction layouts: changing BV can alter BF16 rounding.
+        bv = config.BV
+        _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)](
+            q,
+            k,
+            v,
+            a,
+            b,
+            a_log,
+            bias,
+            self.state[linear_layer_idx],
+            self.keys[linear_layer_idx],
+            self.values[linear_layer_idx],
+            self.decays[linear_layer_idx],
+            self.betas[linear_layer_idx],
+            reqs,
+            cu_seqlens,
+            None,
+            out,
+            q.stride(1),
+            k.stride(1),
+            v.stride(1),
+            a.stride(0),
+            b.stride(0),
+            slots,
+            q.shape[-2],
+            hv,
+            self.num_key_heads,
+            kd,
+            vd,
+            self.verify_width,
+            self.hold_req_id,
+            triton.next_power_of_2(kd),
+            bv,
+            False,
+            self.kda,
+            self.gate_lower_bound,
+            num_warps=config.num_warps,
+            num_stages=config.num_stages,
+        )
+        return out, config
+
+    def accept_updates(self, reqs, accepted_state_indices, run_config=None):
+        """Merge this round's prefix; accepted_state_indices contains per-request last accepted indexes."""
+        assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration"
+        layers, slots, hv, kd, vd = self.state.shape
+        config = run_config or self.run_config
+        if not isinstance(config, SSMRunConfig):
+            config = SSMRunConfig(**config)
+        bv = config.BV
+        _compact[(triton.cdiv(vd, bv), reqs.numel(), layers * hv)](
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            self.state,
+            self.keys,
+            self.values,
+            self.decays,
+            self.betas,
+            reqs,
+            None,
+            accepted_state_indices,
+            None,
+            0,
+            0,
+            0,
+            0,
+            0,
+            slots,
+            hv,
+            hv,
+            self.num_key_heads,
+            kd,
+            vd,
+            self.verify_width,
+            self.hold_req_id,
+            triton.next_power_of_2(kd),
+            bv,
+            True,
+            self.kda,
+            self.gate_lower_bound,
+            num_warps=config.num_warps,
+            num_stages=config.num_stages,
+        )
+
+    def create_commit(self, run_config):
+        if not isinstance(run_config, SSMRunConfig):
+            raise ValueError("compact commit requires the configuration of its forward")
+        return run_config
+
+    def validate_commit(self, context, state_name):
+        if context is None or not context.batches:
+            raise ValueError("compact acceptance requires forward commit batches")
+        for _, config in context.partitions(state_name):
+            if not isinstance(config, SSMRunConfig):
+                raise ValueError(f"missing compact forward configuration for {state_name}")
+
+    def commit(self, reqs, accepted_steps, indices, context, state_name):
+        self.validate_commit(context, state_name)
+        update_accepted_indices(reqs, accepted_steps, indices, context)
+        for partition, config in context.partitions(state_name):
+            self.accept_updates(reqs[partition], indices, run_config=config)