feat: add ReplaySSM state modes for GDN and KDA

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",

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

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",
+]

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

@@ -5,19 +5,33 @@
 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(
     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,
 ) -> 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,
     )

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

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",
+]

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

@@ -5,19 +5,33 @@
 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(
     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,
 ) -> 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,
     )